From c5a2fb42568269c691505ce4660cef27fef5f683 Mon Sep 17 00:00:00 2001 From: Carson Date: Fri, 11 Sep 2026 20:50:22 -0500 Subject: [PATCH 1/9] fix(r): let $server(data_source=) survive a second session safely (#300) $add_table() (which $server(data_source=) calls internally) refuses to run once .server_initialized is TRUE -- a guard meant to stop reconfiguring tables after a conversation is underway, but which also blocked a second Shiny session's own $server(data_source=) call. Even bypassing that, replacing a table cleaned up the resource (and cached query executor) an earlier, still-running session depended on, and the greeting_arg closure read live private$.data_sources/greeter$tables at async-invocation time, so a later session's registration could leak into an earlier session's greeting. Extracts $add_table()'s guard-free core into a private add_or_replace_table(), used directly by $server()'s data_source= path with cleanup_replaced = FALSE so it no longer tears down an earlier session's resource. $server() now also snapshots greeter$tables and passes it (plus the already-per-session data_sources) through mod_server() to greeter$build_client(), so the lazily-invoked greeting reads a point-in-time snapshot instead of live, mutable state. Mirrors the equivalent Python fixes (posit-dev/querychat#302, #303, --- pkg-r/NEWS.md | 2 + pkg-r/R/QueryChat.R | 162 +++++++++++------- pkg-r/R/QueryChatGreeter.R | 18 +- pkg-r/R/querychat_module.R | 9 +- pkg-r/tests/testthat/test-querychat_module.R | 2 +- .../tests/testthat/test-server_data_source.R | 146 ++++++++++++++++ 6 files changed, 273 insertions(+), 66 deletions(-) create mode 100644 pkg-r/tests/testthat/test-server_data_source.R diff --git a/pkg-r/NEWS.md b/pkg-r/NEWS.md index 8f7f9f83e..ed999ae18 100644 --- a/pkg-r/NEWS.md +++ b/pkg-r/NEWS.md @@ -72,6 +72,8 @@ * `$add_table()` no longer rewrites the Shiny module `$id` when the registered table is the only one; the id is now fixed at construction time, matching Python. The rewrite could desync the module namespace from an already-rendered UI when a table was registered between `$ui()` and `$server()` (e.g. via `$server(data_source = )`). (#305) +* `$server(data_source = )` — used for the deferred pattern where a table's data source can only be created inside the server function (e.g. a per-user database connection scoped to Posit Connect OAuth credentials) — only ever worked for the first Shiny session: a second session hit `"Cannot add tables after server initialization"`, and even bypassing that, replacing a table cleaned up the resource an earlier, still-running session depended on, and the auto-generated greeting could reflect the wrong session's table. All three are fixed: each session can now register its own data source, and neither the replaced resource nor the greeting is affected by a later session's registration. (#300) + # querychat 0.3.0 ## New features diff --git a/pkg-r/R/QueryChat.R b/pkg-r/R/QueryChat.R index 8b3b560a3..61ca604ad 100644 --- a/pkg-r/R/QueryChat.R +++ b/pkg-r/R/QueryChat.R @@ -150,6 +150,91 @@ QueryChat <- R6::R6Class( ) }, + # Guard-free core of $add_table(), also called directly by $server()'s + # data_source= path so a session can register its own table even after + # an earlier session's $server() call has 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 (its own query executor may hold a live + # reference to it), so cleaning it up here would pull the resource out + # from under that session. The default (TRUE) preserves $add_table()'s + # existing behavior, where a config-time replacement has exactly one + # owner. + add_or_replace_table = function( + data_source, + table_name, + replace = FALSE, + include_in_greeting = FALSE, + cleanup_replaced = TRUE + ) { + check_bool(include_in_greeting) + check_sql_table_name(table_name) + if (table_name %in% names(private$.data_sources) && !replace) { + cli::cli_abort( + "Table {.val {table_name}} already exists. Use {.code replace = TRUE} to replace." + ) + } + if ( + is_data_source(data_source) && + !identical(data_source$table_name, table_name) + ) { + cli::cli_abort( + c( + "{.arg data_source}'s own table name ({.val {data_source$table_name}}) does not match the given {.arg table_name} ({.val {table_name}}).", + "i" = "Pass a matching {.arg table_name}, or omit it to use {.val {data_source$table_name}}." + ) + ) + } + normalized <- normalize_data_source(data_source, table_name) + + other_sources <- private$.data_sources[ + names(private$.data_sources) != table_name + ] + check_source_compatibility(other_sources, normalized, table_name) + + next_sources <- private$.data_sources + next_sources[[table_name]] <- normalized + + private$auto_fill_data_description(next_sources) + tryCatch( + { + private$build_system_prompt(data_sources = next_sources) + }, + error = function(e) { + if (!inherits(data_source, "DataSource")) { + normalized$cleanup() + } + stop(e) + } + ) + + old_source <- private$.data_sources[[table_name]] + private$.data_sources <- next_sources + if ( + cleanup_replaced && + !is.null(old_source) && + !identical(old_source, normalized) + ) { + old_source$cleanup() + } + + if (!is.null(private$.query_executor)) { + if (cleanup_replaced) { + tryCatch(private$.query_executor$cleanup(), error = function(e) NULL) + } + private$.query_executor <- NULL + } + + # Guard against duplicates: the per-session $server(data_source=) path + # reaches this once per session under the same table name. + if (isTRUE(include_in_greeting) && !(table_name %in% self$greeter$tables)) { + self$greeter$tables <- c(self$greeter$tables, table_name) + } + + invisible(NULL) + }, + create_session_client = function( client_spec = NULL, tools = NA, @@ -434,62 +519,12 @@ QueryChat <- R6::R6Class( if (private$.server_initialized) { cli::cli_abort("Cannot add tables after server initialization.") } - check_bool(include_in_greeting) - check_sql_table_name(table_name) - if (table_name %in% names(private$.data_sources) && !replace) { - cli::cli_abort( - "Table {.val {table_name}} already exists. Use {.code replace = TRUE} to replace." - ) - } - if ( - is_data_source(data_source) && - !identical(data_source$table_name, table_name) - ) { - cli::cli_abort( - c( - "{.arg data_source}'s own table name ({.val {data_source$table_name}}) does not match the given {.arg table_name} ({.val {table_name}}).", - "i" = "Pass a matching {.arg table_name}, or omit it to use {.val {data_source$table_name}}." - ) - ) - } - normalized <- normalize_data_source(data_source, table_name) - - other_sources <- private$.data_sources[ - names(private$.data_sources) != table_name - ] - check_source_compatibility(other_sources, normalized, table_name) - - next_sources <- private$.data_sources - next_sources[[table_name]] <- normalized - - private$auto_fill_data_description(next_sources) - tryCatch( - { - private$build_system_prompt(data_sources = next_sources) - }, - error = function(e) { - if (!inherits(data_source, "DataSource")) { - normalized$cleanup() - } - stop(e) - } + private$add_or_replace_table( + data_source, + table_name, + replace = replace, + include_in_greeting = include_in_greeting ) - - old_source <- private$.data_sources[[table_name]] - private$.data_sources <- next_sources - if (!is.null(old_source) && !identical(old_source, normalized)) { - old_source$cleanup() - } - - if (!is.null(private$.query_executor)) { - tryCatch(private$.query_executor$cleanup(), error = function(e) NULL) - private$.query_executor <- NULL - } - - if (isTRUE(include_in_greeting)) { - self$greeter$tables <- c(self$greeter$tables, table_name) - } - invisible(self) }, @@ -1121,11 +1156,12 @@ QueryChat <- R6::R6Class( ) ) } - self$add_table( + private$add_or_replace_table( data_source, tbl_name, replace = TRUE, - include_in_greeting = TRUE + include_in_greeting = TRUE, + cleanup_replaced = FALSE ) } @@ -1174,7 +1210,8 @@ QueryChat <- R6::R6Class( tools = self$tools, history = resolved_history, greeter = self$greeter, - greeting_base = base_client + greeting_base = base_client, + greeting_tables = self$greeter$tables ) result }, @@ -1217,10 +1254,15 @@ QueryChat <- R6::R6Class( return(invisible(value)) } if (is.null(private$.greeter)) { - client_factory <- function(tables, prompt, base = NULL) { + client_factory <- function( + tables, + prompt, + base = NULL, + data_sources = NULL + ) { sp <- QueryChatSystemPrompt$new( prompt_template = prompt, - data_sources = private$.data_sources, + data_sources = data_sources %||% private$.data_sources, data_description = private$.data_description, extra_instructions = NULL, categorical_threshold = private$.categorical_threshold, diff --git a/pkg-r/R/QueryChatGreeter.R b/pkg-r/R/QueryChatGreeter.R index 1af88edff..9e7103c1d 100644 --- a/pkg-r/R/QueryChatGreeter.R +++ b/pkg-r/R/QueryChatGreeter.R @@ -13,7 +13,7 @@ QueryChatGreeter <- R6::R6Class( .prompt = NULL ), public = list( - #' @param client_factory function(tables, prompt, base) returning a configured greeting client. + #' @param client_factory function(tables, prompt, base, data_sources) returning a configured greeting client. initialize = function(client_factory) { private$.client_factory <- client_factory private$.tables <- character() @@ -26,8 +26,20 @@ QueryChatGreeter <- R6::R6Class( #' @description Build a fresh greeting client (no history) configured with the greeting system prompt. #' @param base Optional resolved client to clone (resolve-once base from `$server()`). - build_client = function(base = NULL) { - private$.client_factory(private$.tables, private$.prompt, base) + #' @param tables Advanced/internal: overrides `$tables` for this call only. + #' Used by `mod_server()` to build the greeting from a point-in-time + #' snapshot captured when a Shiny session's `$server()` call ran, rather + #' than the live (and possibly since-mutated) `$tables`. + #' @param data_sources Advanced/internal: overrides the QueryChat + #' instance's data sources for this call only, for the same reason as + #' `tables`. + build_client = function(base = NULL, tables = NULL, data_sources = NULL) { + private$.client_factory( + tables %||% private$.tables, + private$.prompt, + base, + data_sources = data_sources + ) }, #' @description Generate a greeting synchronously and return it as text. diff --git a/pkg-r/R/querychat_module.R b/pkg-r/R/querychat_module.R index 710f21042..ab8e19fa5 100644 --- a/pkg-r/R/querychat_module.R +++ b/pkg-r/R/querychat_module.R @@ -52,7 +52,8 @@ mod_server <- function( tools, history, greeter = NULL, - greeting_base = NULL + greeting_base = NULL, + greeting_tables = NULL ) { shiny::moduleServer(id, function(input, output, session) { current_table_val <- shiny::reactiveVal(NULL, label = "current_table") @@ -130,7 +131,11 @@ mod_server <- function( "i" = "For faster startup, lower cost, and determinism, consider providing a {.arg greeting} to {.fn QueryChat}.", "i" = "You can use your {.help querychat::QueryChat} object's {.fn $generate_greeting} method to generate a greeting." )) - greeting_client <- greeter$build_client(greeting_base) + greeting_client <- greeter$build_client( + greeting_base, + tables = greeting_tables, + data_sources = data_sources + ) stream <- greeting_client$stream_async(GREETING_PROMPT) shinychat::chat_greeting(stream, persistent = TRUE) } diff --git a/pkg-r/tests/testthat/test-querychat_module.R b/pkg-r/tests/testthat/test-querychat_module.R index 87ef47616..6688d9157 100644 --- a/pkg-r/tests/testthat/test-querychat_module.R +++ b/pkg-r/tests/testthat/test-querychat_module.R @@ -704,7 +704,7 @@ test_that("mod_server() builds the auto-generated greeting from the greeter, not build_client_calls <- list() fake_greeter <- list( - build_client = function(base = NULL) { + build_client = function(base = NULL, tables = NULL, data_sources = NULL) { build_client_calls[[length(build_client_calls) + 1L]] <<- base fake_greeting_client } diff --git a/pkg-r/tests/testthat/test-server_data_source.R b/pkg-r/tests/testthat/test-server_data_source.R new file mode 100644 index 000000000..aaf6b9938 --- /dev/null +++ b/pkg-r/tests/testthat/test-server_data_source.R @@ -0,0 +1,146 @@ +# Tests for QueryChat$server(data_source = ) surviving a second Shiny +# session without corrupting an earlier, still-running session's resources +# or greeting (posit-dev/querychat#300). + +# Any non-NULL value satisfies $server()'s `is.null(session)` guard; the +# value itself is never otherwise used (it isn't threaded into mod_server()). +fake_shiny_session <- function() structure(list(), class = "ShinySession") + +# R6 instances lock existing method bindings, so spying on $cleanup() +# requires unlocking it first. +spy_on_cleanup <- function(obj) { + called <- FALSE + unlockBinding("cleanup", obj) + obj$cleanup <- function() called <<- TRUE + function() called +} + +# Mock mod_server() to capture its call kwargs instead of actually building a +# Shiny module, and return a list() so $server() has something to return. +local_captured_mod_server <- function(env = parent.frame()) { + calls <- new.env(parent = emptyenv()) + calls$args <- list() + + testthat::local_mocked_bindings( + mod_server = function(...) { + calls$args[[length(calls$args) + 1L]] <- list(...) + list() + }, + .package = "querychat", + .env = env + ) + calls +} + +describe("QueryChat$server(data_source=) survives a second session", { + it("a second session's call does not raise", { + skip_if_no_dataframe_engine() + calls <- local_captured_mod_server() + + qc <- QueryChat$new(NULL, "users", greeting = "Test") + withr::defer(qc$cleanup()) + + qc$server(data_source = new_users_df(), session = fake_shiny_session()) + expect_no_error( + qc$server(data_source = new_users_df(), session = fake_shiny_session()) + ) + + expect_equal(names(calls$args[[2]]$data_sources), "users") + }) + + it("the public $add_table() guard is still enforced", { + skip_if_no_dataframe_engine() + local_captured_mod_server() + + qc <- QueryChat$new(NULL, "users", greeting = "Test") + withr::defer(qc$cleanup()) + + qc$server(data_source = new_users_df(), session = fake_shiny_session()) + + expect_error( + qc$add_table(new_users_df(), "other"), + "Cannot add tables after server initialization" + ) + }) +}) + +describe("QueryChat$server(data_source=) cleanup safety", { + it("a second session's call does not clean up the first session's source", { + skip_if_no_dataframe_engine() + local_captured_mod_server() + + qc <- QueryChat$new(NULL, "users", greeting = "Test") + withr::defer(qc$cleanup()) + + qc$server(data_source = new_users_df(), session = fake_shiny_session()) + first_source <- qc_data_source(qc, "users") + cleaned_up <- spy_on_cleanup(first_source) + + qc$server(data_source = new_users_df(), session = fake_shiny_session()) + + expect_false(cleaned_up()) + }) + + it("the public add_table(replace=TRUE) still cleans up the old source", { + skip_if_no_dataframe_engine() + + qc <- QueryChat$new(new_users_df(), "users", greeting = "Test") + withr::defer(qc$cleanup()) + + first_source <- qc_data_source(qc, "users") + cleaned_up <- spy_on_cleanup(first_source) + + qc$add_table(new_users_df(), "users", replace = TRUE) + + expect_true(cleaned_up()) + }) + + it("a second session's call does not clean up the first session's query executor", { + skip_if_no_dataframe_engine() + local_captured_mod_server() + + qc <- QueryChat$new(NULL, "users", greeting = "Test") + withr::defer(qc$cleanup()) + + qc$server(data_source = new_users_df(), session = fake_shiny_session()) + first_executor <- qc$.__enclos_env__$private$.query_executor + cleaned_up <- spy_on_cleanup(first_executor) + + qc$server(data_source = new_users_df(), session = fake_shiny_session()) + + expect_false(cleaned_up()) + }) + + it("the public add_table(replace=TRUE) still cleans up the old query executor", { + skip_if_no_dataframe_engine() + + qc <- QueryChat$new(new_users_df(), "users", greeting = "Test") + withr::defer(qc$cleanup()) + + # Force the executor to be built, mirroring a session that's already + # queried through it before a config-time replace happens. + qc$.__enclos_env__$private$.query_executor <- build_query_executor( + qc$.__enclos_env__$private$.data_sources + ) + first_executor <- qc$.__enclos_env__$private$.query_executor + cleaned_up <- spy_on_cleanup(first_executor) + + qc$add_table(new_users_df(), "users", replace = TRUE) + + expect_true(cleaned_up()) + }) +}) + +describe("QueryChat$server(data_source=) greeting snapshot", { + it("passes a greeting_tables snapshot to mod_server", { + skip_if_no_dataframe_engine() + calls <- local_captured_mod_server() + + qc <- QueryChat$new(NULL, "users", greeting = "Test") + withr::defer(qc$cleanup()) + + qc$server(data_source = new_users_df(), session = fake_shiny_session()) + + expect_equal(calls$args[[1]]$greeting_tables, "users") + }) +}) From 13d0904e23777db4c45f232fb416bc270d60ba72 Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 10:03:40 -0500 Subject: [PATCH 2/9] fix(r): don't duplicate greeter$tables on per-session registration add_or_replace_table() appended table_name to greeter$tables unconditionally, so each session's $server(data_source=) call added another copy of the same name, and later sessions' greeting snapshots would list the table repeatedly. Match Python's dedup guard. --- pkg-r/R/QueryChat.R | 4 +++- pkg-r/tests/testthat/test-server_data_source.R | 14 ++++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/pkg-r/R/QueryChat.R b/pkg-r/R/QueryChat.R index 61ca604ad..97c3ec8a0 100644 --- a/pkg-r/R/QueryChat.R +++ b/pkg-r/R/QueryChat.R @@ -228,7 +228,9 @@ QueryChat <- R6::R6Class( # Guard against duplicates: the per-session $server(data_source=) path # reaches this once per session under the same table name. - if (isTRUE(include_in_greeting) && !(table_name %in% self$greeter$tables)) { + if ( + isTRUE(include_in_greeting) && !(table_name %in% self$greeter$tables) + ) { self$greeter$tables <- c(self$greeter$tables, table_name) } diff --git a/pkg-r/tests/testthat/test-server_data_source.R b/pkg-r/tests/testthat/test-server_data_source.R index aaf6b9938..22ec5de86 100644 --- a/pkg-r/tests/testthat/test-server_data_source.R +++ b/pkg-r/tests/testthat/test-server_data_source.R @@ -143,4 +143,18 @@ describe("QueryChat$server(data_source=) greeting snapshot", { expect_equal(calls$args[[1]]$greeting_tables, "users") }) + + it("per-session registration does not duplicate greeter$tables", { + skip_if_no_dataframe_engine() + calls <- local_captured_mod_server() + + qc <- QueryChat$new(NULL, "users", greeting = "Test") + withr::defer(qc$cleanup()) + + qc$server(data_source = new_users_df(), session = fake_shiny_session()) + qc$server(data_source = new_users_df(), session = fake_shiny_session()) + + expect_equal(qc$greeter$tables, "users") + expect_equal(calls$args[[2]]$greeting_tables, "users") + }) }) From 18fe7f76e8aece85db5dab31762f7c9952e9f506 Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 10:18:44 -0500 Subject: [PATCH 3/9] docs(r): trim comments to what isn't inferable from the code --- pkg-r/R/QueryChat.R | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/pkg-r/R/QueryChat.R b/pkg-r/R/QueryChat.R index 97c3ec8a0..cf4a3001e 100644 --- a/pkg-r/R/QueryChat.R +++ b/pkg-r/R/QueryChat.R @@ -151,16 +151,10 @@ QueryChat <- R6::R6Class( }, # Guard-free core of $add_table(), also called directly by $server()'s - # data_source= path so a session can register its own table even after - # an earlier session's $server() call has 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 (its own query executor may hold a live - # reference to it), so cleaning it up here would pull the resource out - # from under that session. The default (TRUE) preserves $add_table()'s - # existing behavior, where a config-time replacement has exactly one - # owner. + # per-session data_source= path (which must work even after an earlier + # session set .server_initialized). cleanup_replaced = FALSE is for that + # path: the replaced table may still be in use by an earlier session, so + # cleanup becomes the caller's responsibility. add_or_replace_table = function( data_source, table_name, @@ -474,8 +468,7 @@ QueryChat <- R6::R6Class( self$greeter$tables <- c(self$greeter$tables, normalized$table_name) self$id <- id %||% sprintf("querychat_%s", normalized$table_name) } else { - # Deferred pattern: data_source is NULL. table_name is optional here; - # explicit NULL is treated the same as omitting it. + # An explicit table_name = NULL is treated the same as omitting it. table_name_given <- !is_missing(table_name) && !is.null(table_name) if (table_name_given) { private$.deferred_table_name <- table_name From 43971d0f463bf36820fdc92c7d88f892bd7f6253 Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 11:01:28 -0500 Subject: [PATCH 4/9] test(r): lock in mixed config-time add_table() + per-session server(data_source=) behavior Covers: same-name (or 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, so a later session's snapshot includes an earlier session's differently-named table. --- .../tests/testthat/test-server_data_source.R | 89 +++++++++++++++++++ 1 file changed, 89 insertions(+) diff --git a/pkg-r/tests/testthat/test-server_data_source.R b/pkg-r/tests/testthat/test-server_data_source.R index 22ec5de86..f23a8d21e 100644 --- a/pkg-r/tests/testthat/test-server_data_source.R +++ b/pkg-r/tests/testthat/test-server_data_source.R @@ -158,3 +158,92 @@ describe("QueryChat$server(data_source=) greeting snapshot", { expect_equal(calls$args[[2]]$greeting_tables, "users") }) }) + +describe("Mixing config-time $add_table() with $server(data_source=)", { + it("server(data_source=) without a name replaces the config-time table", { + skip_if_no_dataframe_engine() + calls <- local_captured_mod_server() + + qc <- QueryChat$new(NULL, greeting = "Test") + withr::defer(qc$cleanup()) + + config_df <- data.frame(id = 1:3) + session_df <- data.frame(id = 4:6) + qc$add_table(config_df, "orders") + + qc$server(data_source = session_df, session = fake_shiny_session()) + + # Same table name, but the session's data replaces the config-time data + expect_equal(names(calls$args[[1]]$data_sources), "orders") + expect_equal(calls$args[[1]]$data_sources$orders$get_data()$id, 4:6) + }) + + it("replacing a config-time table does not clean it up", { + skip_if_no_dataframe_engine() + local_captured_mod_server() + + qc <- QueryChat$new(NULL, greeting = "Test") + withr::defer(qc$cleanup()) + + qc$add_table(new_users_df(), "orders") + config_source <- qc_data_source(qc, "orders") + cleaned_up <- spy_on_cleanup(config_source) + + qc$server(data_source = new_users_df(), session = fake_shiny_session()) + + # Consistent with per-session replacement: the replaced source's cleanup + # is left to whoever created it + expect_false(cleaned_up()) + }) + + it("server(data_source=, table_name=) adds a second table alongside the config-time one", { + skip_if_no_dataframe_engine() + calls <- local_captured_mod_server() + + qc <- QueryChat$new(NULL, greeting = "Test") + withr::defer(qc$cleanup()) + + config_df <- data.frame(id = 1:3) + session_df <- data.frame(id = 4:6) + qc$add_table(config_df, "orders") + + qc$server( + data_source = session_df, + table_name = "returns", + session = fake_shiny_session() + ) + + expect_equal(names(calls$args[[1]]$data_sources), c("orders", "returns")) + # The config-time table's own data is untouched + expect_equal(calls$args[[1]]$data_sources$orders$get_data()$id, 1:3) + expect_equal(calls$args[[1]]$data_sources$returns$get_data()$id, 4:6) + }) + + it("registration state is shared: a later session's snapshot includes an earlier session's differently-named table", { + skip_if_no_dataframe_engine() + calls <- local_captured_mod_server() + + qc <- QueryChat$new(NULL, greeting = "Test") + withr::defer(qc$cleanup()) + + qc$add_table(data.frame(id = 1:3), "orders") + + # Session 1 adds its own table alongside the config-time one + qc$server( + data_source = data.frame(id = 4:6), + table_name = "returns", + session = fake_shiny_session() + ) + # Session 2 replaces "orders" only -- but still sees session 1's table, + # since the registry is shared and cumulative across sessions + qc$server( + data_source = data.frame(id = 7:9), + table_name = "orders", + session = fake_shiny_session() + ) + + expect_equal(names(calls$args[[2]]$data_sources), c("orders", "returns")) + expect_equal(calls$args[[2]]$data_sources$orders$get_data()$id, 7:9) + expect_equal(calls$args[[2]]$data_sources$returns$get_data()$id, 4:6) + }) +}) From bbfe168018c6d826977a4c02b535b775cd33274e Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 11:57:22 -0500 Subject: [PATCH 5/9] fix(r): scope $server(data_source=) cleanup to live sessions; snapshot greeting data_description - Track live sessions on the instance (.active_sessions, incremented per $server() call, decremented via session$onSessionEnded()) and key shared-state decisions off it: $server(data_source=) cleans up the resource it replaces exactly when no session is live (cleanup_replaced = .active_sessions == 0), and the $add_table()/$add_tables()/$remove_table() guards block only while a session is active. A resource registered by a session that has since ended is reclaimed when replaced; a missed onSessionEnded degrades to over-conservative cleanup skipping, never premature cleanup. - The greeting snapshot now also captures data_description, so a later session's re-inferred description can't leak into an earlier session's lazily-generated greeting. - The mod_server() greeting test now asserts the snapshot args (tables, data_sources, data_description) are actually forwarded to greeter$build_client(), not just the base client. --- pkg-r/NEWS.md | 2 +- pkg-r/R/QueryChat.R | 37 ++-- pkg-r/R/QueryChatGreeter.R | 13 +- pkg-r/R/querychat_module.R | 6 +- pkg-r/tests/testthat/test-QueryChat.R | 6 +- pkg-r/tests/testthat/test-querychat_module.R | 30 +++- .../tests/testthat/test-server_data_source.R | 165 +++++++++++++++++- 7 files changed, 226 insertions(+), 33 deletions(-) diff --git a/pkg-r/NEWS.md b/pkg-r/NEWS.md index ed999ae18..910774153 100644 --- a/pkg-r/NEWS.md +++ b/pkg-r/NEWS.md @@ -72,7 +72,7 @@ * `$add_table()` no longer rewrites the Shiny module `$id` when the registered table is the only one; the id is now fixed at construction time, matching Python. The rewrite could desync the module namespace from an already-rendered UI when a table was registered between `$ui()` and `$server()` (e.g. via `$server(data_source = )`). (#305) -* `$server(data_source = )` — used for the deferred pattern where a table's data source can only be created inside the server function (e.g. a per-user database connection scoped to Posit Connect OAuth credentials) — only ever worked for the first Shiny session: a second session hit `"Cannot add tables after server initialization"`, and even bypassing that, replacing a table cleaned up the resource an earlier, still-running session depended on, and the auto-generated greeting could reflect the wrong session's table. All three are fixed: each session can now register its own data source, and neither the replaced resource nor the greeting is affected by a later session's registration. (#300) +* `$server(data_source = )` — used for the deferred pattern where a table's data source can only be created inside the server function (e.g. a per-user database connection scoped to Posit Connect OAuth credentials) — only ever worked for the first Shiny session: a second session hit `"Cannot add tables after server initialization"`, and even bypassing that, replacing a table cleaned up the resource an earlier, still-running session depended on, and the auto-generated greeting could reflect the wrong session's table or inferred data description. All three are fixed: each session can now register its own data source, and neither the replaced resource nor the greeting is affected by a later session's registration. Resource cleanup and the `$add_table()`/`$remove_table()` guards now key off *live* sessions (tracked via `session$onSessionEnded()`), so a resource registered by a session that has since ended is reclaimed when replaced. (#300) # querychat 0.3.0 diff --git a/pkg-r/R/QueryChat.R b/pkg-r/R/QueryChat.R index cf4a3001e..e1fcbc452 100644 --- a/pkg-r/R/QueryChat.R +++ b/pkg-r/R/QueryChat.R @@ -92,7 +92,9 @@ QueryChat <- R6::R6Class( .data_sources = list(), .deferred_table_name = NULL, .query_executor = NULL, - .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. + .active_sessions = 0L, .client_spec = NULL, .client_console = NULL, .system_prompt = NULL, @@ -151,8 +153,8 @@ QueryChat <- R6::R6Class( }, # Guard-free core of $add_table(), also called directly by $server()'s - # per-session data_source= path (which must work even after an earlier - # session set .server_initialized). cleanup_replaced = FALSE is for that + # per-session data_source= path (which must work even while earlier + # sessions are still running). cleanup_replaced = FALSE is for that # path: the replaced table may still be in use by an earlier session, so # cleanup becomes the caller's responsibility. add_or_replace_table = function( @@ -511,8 +513,8 @@ QueryChat <- R6::R6Class( replace = FALSE, include_in_greeting = FALSE ) { - if (private$.server_initialized) { - cli::cli_abort("Cannot add tables after server initialization.") + if (private$.active_sessions > 0) { + cli::cli_abort("Cannot add tables while a server session is active.") } private$add_or_replace_table( data_source, @@ -548,8 +550,8 @@ QueryChat <- R6::R6Class( replace = FALSE, include_in_greeting = FALSE ) { - if (private$.server_initialized) { - cli::cli_abort("Cannot add tables after server initialization.") + if (private$.active_sessions > 0) { + cli::cli_abort("Cannot add tables while a server session is active.") } if (!inherits(conn, "DBIConnection")) { cli::cli_abort( @@ -644,8 +646,8 @@ QueryChat <- R6::R6Class( #' #' @return Invisibly returns `self` for chaining. remove_table = function(table_name) { - if (private$.server_initialized) { - cli::cli_abort("Cannot remove tables after server initialization.") + if (private$.active_sessions > 0) { + cli::cli_abort("Cannot remove tables while a server session is active.") } if (!table_name %in% names(private$.data_sources)) { cli::cli_abort("Table {.val {table_name}} not found.") @@ -1156,13 +1158,18 @@ QueryChat <- R6::R6Class( tbl_name, replace = TRUE, include_in_greeting = TRUE, - cleanup_replaced = FALSE + # A live session may still be using the replaced source, so only + # clean it up once no sessions are active. + cleanup_replaced = private$.active_sessions == 0 ) } private$require_initialized("$server") - private$.server_initialized <- TRUE + private$.active_sessions <- private$.active_sessions + 1L + session$onSessionEnded(function() { + private$.active_sessions <- private$.active_sessions - 1L + }) if (is.null(private$.query_executor)) { private$.query_executor <- build_query_executor(private$.data_sources) @@ -1206,7 +1213,8 @@ QueryChat <- R6::R6Class( history = resolved_history, greeter = self$greeter, greeting_base = base_client, - greeting_tables = self$greeter$tables + greeting_tables = self$greeter$tables, + greeting_data_description = private$.data_description ) result }, @@ -1253,12 +1261,13 @@ QueryChat <- R6::R6Class( tables, prompt, base = NULL, - data_sources = NULL + data_sources = NULL, + data_description = NULL ) { sp <- QueryChatSystemPrompt$new( prompt_template = prompt, data_sources = data_sources %||% private$.data_sources, - data_description = private$.data_description, + data_description = data_description %||% private$.data_description, extra_instructions = NULL, categorical_threshold = private$.categorical_threshold, data_dicts = private$.data_dicts, diff --git a/pkg-r/R/QueryChatGreeter.R b/pkg-r/R/QueryChatGreeter.R index 9e7103c1d..cd9b0557d 100644 --- a/pkg-r/R/QueryChatGreeter.R +++ b/pkg-r/R/QueryChatGreeter.R @@ -33,12 +33,21 @@ QueryChatGreeter <- R6::R6Class( #' @param data_sources Advanced/internal: overrides the QueryChat #' instance's data sources for this call only, for the same reason as #' `tables`. - build_client = function(base = NULL, tables = NULL, data_sources = NULL) { + #' @param data_description Advanced/internal: overrides the QueryChat + #' instance's inferred data description for this call only, for the same + #' reason as `tables`. + build_client = function( + base = NULL, + tables = NULL, + data_sources = NULL, + data_description = NULL + ) { private$.client_factory( tables %||% private$.tables, private$.prompt, base, - data_sources = data_sources + data_sources = data_sources, + data_description = data_description ) }, diff --git a/pkg-r/R/querychat_module.R b/pkg-r/R/querychat_module.R index ab8e19fa5..780bcc5ed 100644 --- a/pkg-r/R/querychat_module.R +++ b/pkg-r/R/querychat_module.R @@ -53,7 +53,8 @@ mod_server <- function( history, greeter = NULL, greeting_base = NULL, - greeting_tables = NULL + greeting_tables = NULL, + greeting_data_description = NULL ) { shiny::moduleServer(id, function(input, output, session) { current_table_val <- shiny::reactiveVal(NULL, label = "current_table") @@ -134,7 +135,8 @@ mod_server <- function( greeting_client <- greeter$build_client( greeting_base, tables = greeting_tables, - data_sources = data_sources + data_sources = data_sources, + data_description = greeting_data_description ) stream <- greeting_client$stream_async(GREETING_PROMPT) shinychat::chat_greeting(stream, persistent = TRUE) diff --git a/pkg-r/tests/testthat/test-QueryChat.R b/pkg-r/tests/testthat/test-QueryChat.R index 0b810131c..ae841dfaf 100644 --- a/pkg-r/tests/testthat/test-QueryChat.R +++ b/pkg-r/tests/testthat/test-QueryChat.R @@ -1340,13 +1340,13 @@ describe("QueryChat$add_tables()", { ) }) - it("calling after server initialization raises error", { + it("calling while a server session is active raises error", { conn <- local_multi_table_conn() qc <- QueryChat$new(NULL, "placeholder", greeting = "Test") - qc$.__enclos_env__$private$.server_initialized <- TRUE + qc$.__enclos_env__$private$.active_sessions <- 1L expect_error( qc$add_tables(conn), - "after server initialization" + "while a server session is active" ) }) diff --git a/pkg-r/tests/testthat/test-querychat_module.R b/pkg-r/tests/testthat/test-querychat_module.R index 6688d9157..464fa80c7 100644 --- a/pkg-r/tests/testthat/test-querychat_module.R +++ b/pkg-r/tests/testthat/test-querychat_module.R @@ -704,8 +704,18 @@ test_that("mod_server() builds the auto-generated greeting from the greeter, not build_client_calls <- list() fake_greeter <- list( - build_client = function(base = NULL, tables = NULL, data_sources = NULL) { - build_client_calls[[length(build_client_calls) + 1L]] <<- base + build_client = function( + base = NULL, + tables = NULL, + data_sources = NULL, + data_description = NULL + ) { + build_client_calls[[length(build_client_calls) + 1L]] <<- list( + base = base, + tables = tables, + data_sources = data_sources, + data_description = data_description + ) fake_greeting_client } ) @@ -732,6 +742,8 @@ test_that("mod_server() builds the auto-generated greeting from the greeter, not tools = "query", greeter = fake_greeter, greeting_base = "base-client", + greeting_tables = "test_table", + greeting_data_description = "snapshot description", history = TRUE ), { @@ -742,9 +754,21 @@ test_that("mod_server() builds the auto-generated greeting from the greeter, not # dedicated greeting prompt), not from a second call to the main client # factory with tools = NULL (which would use the query system prompt). expect_equal(length(build_client_calls), 1L) - expect_equal(build_client_calls[[1]], "base-client") + expect_equal(build_client_calls[[1]]$base, "base-client") expect_equal(length(main_client_calls), 1L) expect_false(is.null(greeting_stream_prompt)) + + # The greeting must be built from the session's point-in-time snapshot, + # not live QueryChat state that a later session may have mutated + expect_identical(build_client_calls[[1]]$tables, "test_table") + expect_identical( + build_client_calls[[1]]$data_sources, + list(test_table = ds) + ) + expect_identical( + build_client_calls[[1]]$data_description, + "snapshot description" + ) } ) }) diff --git a/pkg-r/tests/testthat/test-server_data_source.R b/pkg-r/tests/testthat/test-server_data_source.R index f23a8d21e..946525644 100644 --- a/pkg-r/tests/testthat/test-server_data_source.R +++ b/pkg-r/tests/testthat/test-server_data_source.R @@ -2,9 +2,27 @@ # session without corrupting an earlier, still-running session's resources # or greeting (posit-dev/querychat#300). -# Any non-NULL value satisfies $server()'s `is.null(session)` guard; the -# value itself is never otherwise used (it isn't threaded into mod_server()). -fake_shiny_session <- function() structure(list(), class = "ShinySession") +# A fake session records onSessionEnded() callbacks so tests can simulate the +# session ending via $end(). The session isn't threaded into mod_server(); +# $server() otherwise only NULL-checks it and registers session-end callbacks. +fake_shiny_session <- function() { + ended_callbacks <- list() + structure( + list( + onSessionEnded = function(cb) { + ended_callbacks[[length(ended_callbacks) + 1L]] <<- cb + invisible() + }, + end = function() { + for (cb in ended_callbacks) { + cb() + } + invisible() + } + ), + class = "ShinySession" + ) +} # R6 instances lock existing method bindings, so spying on $cleanup() # requires unlocking it first. @@ -59,7 +77,7 @@ describe("QueryChat$server(data_source=) survives a second session", { expect_error( qc$add_table(new_users_df(), "other"), - "Cannot add tables after server initialization" + "Cannot add tables while a server session is active" ) }) }) @@ -129,6 +147,98 @@ describe("QueryChat$server(data_source=) cleanup safety", { expect_true(cleaned_up()) }) + + it("a first $server(data_source=) call still cleans up the replaced constructor-registered source", { + skip_if_no_dataframe_engine() + local_captured_mod_server() + + qc <- QueryChat$new(new_users_df(), "users", greeting = "Test") + withr::defer(qc$cleanup()) + + constructor_source <- qc_data_source(qc, "users") + cleaned_up <- spy_on_cleanup(constructor_source) + + qc$server(data_source = new_users_df(), session = fake_shiny_session()) + + # No session can still be using it, so cleanup-on-replace holds here + expect_true(cleaned_up()) + }) + + it("a second session's call skips cleanup even when the first call cleaned up", { + skip_if_no_dataframe_engine() + local_captured_mod_server() + + qc <- QueryChat$new(new_users_df(), "users", greeting = "Test") + withr::defer(qc$cleanup()) + + qc$server(data_source = new_users_df(), session = fake_shiny_session()) + session1_source <- qc_data_source(qc, "users") + cleaned_up <- spy_on_cleanup(session1_source) + + qc$server(data_source = new_users_df(), session = fake_shiny_session()) + + expect_false(cleaned_up()) + }) +}) + +describe("QueryChat$server(data_source=) session lifecycle", { + # The hazard behind cleanup-on-replace and the add/remove_table guards is + # *live* sessions, not past ones: an ended session can no longer be using + # a resource it registered. + + it("a replaced source is cleaned up once the session that registered it has ended", { + skip_if_no_dataframe_engine() + local_captured_mod_server() + + qc <- QueryChat$new(NULL, "users", greeting = "Test") + withr::defer(qc$cleanup()) + + session1 <- fake_shiny_session() + qc$server(data_source = new_users_df(), session = session1) + first_source <- qc_data_source(qc, "users") + cleaned_up <- spy_on_cleanup(first_source) + + session1$end() + + qc$server(data_source = new_users_df(), session = fake_shiny_session()) + + expect_true(cleaned_up()) + }) + + it("a replaced source survives while any session is live", { + skip_if_no_dataframe_engine() + local_captured_mod_server() + + qc <- QueryChat$new(NULL, "users", greeting = "Test") + withr::defer(qc$cleanup()) + + session1 <- fake_shiny_session() + qc$server(data_source = new_users_df(), session = session1) + qc$server(data_source = new_users_df(), session = fake_shiny_session()) + second_source <- qc_data_source(qc, "users") + cleaned_up <- spy_on_cleanup(second_source) + + session1$end() # session 1 ends; session 2 still live + + qc$server(data_source = new_users_df(), session = fake_shiny_session()) + + expect_false(cleaned_up()) + }) + + it("$add_table() is allowed once all sessions have ended", { + skip_if_no_dataframe_engine() + local_captured_mod_server() + + qc <- QueryChat$new(NULL, "users", greeting = "Test") + withr::defer(qc$cleanup()) + + session1 <- fake_shiny_session() + qc$server(data_source = new_users_df(), session = session1) + + session1$end() + + expect_no_error(qc$add_table(new_users_df(), "other")) + }) }) describe("QueryChat$server(data_source=) greeting snapshot", { @@ -157,6 +267,45 @@ describe("QueryChat$server(data_source=) greeting snapshot", { expect_equal(qc$greeter$tables, "users") expect_equal(calls$args[[2]]$greeting_tables, "users") }) + + it("passes a greeting_data_description snapshot to mod_server", { + skip_if_no_dataframe_engine() + calls <- local_captured_mod_server() + + qc <- QueryChat$new( + NULL, + "users", + greeting = "Test", + data_description = "User accounts" + ) + withr::defer(qc$cleanup()) + + qc$server(data_source = new_users_df(), session = fake_shiny_session()) + + expect_equal(calls$args[[1]]$greeting_data_description, "User accounts") + }) + + it("greeter$build_client() renders a data_description snapshot instead of live state", { + skip_if_no_dataframe_engine() + + qc <- QueryChat$new( + new_users_df(), + "users", + greeting = "Test", + data_description = "live description", + client = mock_ellmer_chat_client() + ) + withr::defer(qc$cleanup()) + + live <- qc$greeter$build_client()$get_system_prompt() + snapshot <- qc$greeter$build_client( + data_description = "snapshot description" + )$get_system_prompt() + + expect_match(live, "live description", fixed = TRUE) + expect_match(snapshot, "snapshot description", fixed = TRUE) + expect_no_match(snapshot, "live description", fixed = TRUE) + }) }) describe("Mixing config-time $add_table() with $server(data_source=)", { @@ -178,7 +327,7 @@ describe("Mixing config-time $add_table() with $server(data_source=)", { expect_equal(calls$args[[1]]$data_sources$orders$get_data()$id, 4:6) }) - it("replacing a config-time table does not clean it up", { + it("replacing a config-time table on the first $server() call cleans it up", { skip_if_no_dataframe_engine() local_captured_mod_server() @@ -191,9 +340,9 @@ describe("Mixing config-time $add_table() with $server(data_source=)", { qc$server(data_source = new_users_df(), session = fake_shiny_session()) - # Consistent with per-session replacement: the replaced source's cleanup - # is left to whoever created it - expect_false(cleaned_up()) + # No session is running yet, so the replaced source has a single owner + # and cleanup-on-replace still holds (only later sessions skip it) + expect_true(cleaned_up()) }) it("server(data_source=, table_name=) adds a second table alongside the config-time one", { From d3ad93bf21c045b7c0624737f18a0f95d43b3b49 Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 12:10:53 -0500 Subject: [PATCH 6/9] fix(r): defer retired resource cleanup; harden greeting snapshot and deferred name validation Address Copilot review on #306: - Retain sources/executors replaced while sessions are live (.retired_resources) and clean them up when the last live session ends (or $cleanup() runs), fixing a leaked executor/connection when cleanup_replaced was FALSE. Mirrors the Python side (#308). - greeter$build_client() distinguishes an omitted data_sources/ data_description override (falls back to live state) from an explicit snapshot, which may itself be NULL -- a session whose snapshot had no inferred description no longer picks up a later session's. - QueryChat$new(NULL, table_name = ...) validates the deferred name at construction instead of failing later at $server() registration. --- pkg-r/NEWS.md | 2 +- pkg-r/R/QueryChat.R | 66 +++++++++++++---- pkg-r/R/QueryChatGreeter.R | 23 ++++-- .../tests/testthat/test-server_data_source.R | 73 +++++++++++++++++++ 4 files changed, 144 insertions(+), 20 deletions(-) diff --git a/pkg-r/NEWS.md b/pkg-r/NEWS.md index 910774153..17159acf2 100644 --- a/pkg-r/NEWS.md +++ b/pkg-r/NEWS.md @@ -72,7 +72,7 @@ * `$add_table()` no longer rewrites the Shiny module `$id` when the registered table is the only one; the id is now fixed at construction time, matching Python. The rewrite could desync the module namespace from an already-rendered UI when a table was registered between `$ui()` and `$server()` (e.g. via `$server(data_source = )`). (#305) -* `$server(data_source = )` — used for the deferred pattern where a table's data source can only be created inside the server function (e.g. a per-user database connection scoped to Posit Connect OAuth credentials) — only ever worked for the first Shiny session: a second session hit `"Cannot add tables after server initialization"`, and even bypassing that, replacing a table cleaned up the resource an earlier, still-running session depended on, and the auto-generated greeting could reflect the wrong session's table or inferred data description. All three are fixed: each session can now register its own data source, and neither the replaced resource nor the greeting is affected by a later session's registration. Resource cleanup and the `$add_table()`/`$remove_table()` guards now key off *live* sessions (tracked via `session$onSessionEnded()`), so a resource registered by a session that has since ended is reclaimed when replaced. (#300) +* `$server(data_source = )` — used for the deferred pattern where a table's data source can only be created inside the server function (e.g. a per-user database connection scoped to Posit Connect OAuth credentials) — only ever worked for the first Shiny session: a second session hit `"Cannot add tables after server initialization"`, and even bypassing that, replacing a table cleaned up the resource an earlier, still-running session depended on, and the auto-generated greeting could reflect the wrong session's table or inferred data description. All three are fixed: each session can now register its own data source, and neither the replaced resource nor the greeting is affected by a later session's registration. Resource cleanup and the `$add_table()`/`$remove_table()` guards now key off *live* sessions (tracked via `session$onSessionEnded()`): a resource registered by a session that has since ended is reclaimed when replaced, and a resource replaced while sessions are still live is retained and cleaned up once the last live session ends (or `$cleanup()` runs). (#300) # querychat 0.3.0 diff --git a/pkg-r/R/QueryChat.R b/pkg-r/R/QueryChat.R index e1fcbc452..2fd033d66 100644 --- a/pkg-r/R/QueryChat.R +++ b/pkg-r/R/QueryChat.R @@ -95,6 +95,10 @@ QueryChat <- R6::R6Class( # Live Shiny session count. Shared-state guards key off this: an ended # session can no longer be using a resource it registered. .active_sessions = 0L, + # 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. + .retired_resources = list(), .client_spec = NULL, .client_console = NULL, .system_prompt = NULL, @@ -156,7 +160,8 @@ QueryChat <- R6::R6Class( # per-session data_source= path (which must work even while earlier # sessions are still running). cleanup_replaced = FALSE is for that # path: the replaced table may still be in use by an earlier session, so - # cleanup becomes the caller's responsibility. + # the retired source and cached executor are retained and cleaned up + # once the last live session ends (see $server()). add_or_replace_table = function( data_source, table_name, @@ -207,17 +212,25 @@ QueryChat <- R6::R6Class( old_source <- private$.data_sources[[table_name]] private$.data_sources <- next_sources - if ( - cleanup_replaced && - !is.null(old_source) && - !identical(old_source, normalized) - ) { - old_source$cleanup() + if (!is.null(old_source) && !identical(old_source, normalized)) { + if (cleanup_replaced) { + old_source$cleanup() + } else { + private$.retired_resources <- c( + private$.retired_resources, + list(old_source) + ) + } } if (!is.null(private$.query_executor)) { if (cleanup_replaced) { tryCatch(private$.query_executor$cleanup(), error = function(e) NULL) + } else { + private$.retired_resources <- c( + private$.retired_resources, + list(private$.query_executor) + ) } private$.query_executor <- NULL } @@ -233,6 +246,18 @@ QueryChat <- R6::R6Class( invisible(NULL) }, + # Clean up resources retired while sessions were still live. Retired + # resources may still be in use by a live session, so this only runs + # once no sessions remain (or from $cleanup()). + flush_retired_resources = function() { + retired <- private$.retired_resources + private$.retired_resources <- list() + for (resource in retired) { + # Best-effort: one failing cleanup must not leave the rest open. + tryCatch(resource$cleanup(), error = function(e) NULL) + } + }, + create_session_client = function( client_spec = NULL, tools = NA, @@ -473,6 +498,9 @@ QueryChat <- R6::R6Class( # An explicit table_name = NULL is treated the same as omitting it. table_name_given <- !is_missing(table_name) && !is.null(table_name) if (table_name_given) { + # Validate now: a bad deferred name would otherwise only surface at + # $server() registration time, after the module id is built. + check_sql_table_name(table_name) private$.deferred_table_name <- table_name } default_id <- if (table_name_given) { @@ -1158,8 +1186,8 @@ QueryChat <- R6::R6Class( tbl_name, replace = TRUE, include_in_greeting = TRUE, - # A live session may still be using the replaced source, so only - # clean it up once no sessions are active. + # A live session may still be using the replaced source, so defer + # its cleanup until no sessions are active. cleanup_replaced = private$.active_sessions == 0 ) } @@ -1169,6 +1197,9 @@ QueryChat <- R6::R6Class( private$.active_sessions <- private$.active_sessions + 1L session$onSessionEnded(function() { private$.active_sessions <- private$.active_sessions - 1L + if (private$.active_sessions == 0) { + private$flush_retired_resources() + } }) if (is.null(private$.query_executor)) { @@ -1243,6 +1274,7 @@ QueryChat <- R6::R6Class( for (source in private$.data_sources) { source$cleanup() } + private$flush_retired_resources() invisible(NULL) } ), @@ -1261,13 +1293,21 @@ QueryChat <- R6::R6Class( tables, prompt, base = NULL, - data_sources = NULL, - data_description = NULL + data_sources, + data_description ) { + # An explicit snapshot (possibly NULL) wins over live state; only + # an omitted override falls back to it. + if (missing(data_sources)) { + data_sources <- private$.data_sources + } + if (missing(data_description)) { + data_description <- private$.data_description + } sp <- QueryChatSystemPrompt$new( prompt_template = prompt, - data_sources = data_sources %||% private$.data_sources, - data_description = data_description %||% private$.data_description, + data_sources = data_sources, + data_description = data_description, extra_instructions = NULL, categorical_threshold = private$.categorical_threshold, data_dicts = private$.data_dicts, diff --git a/pkg-r/R/QueryChatGreeter.R b/pkg-r/R/QueryChatGreeter.R index cd9b0557d..329c57d73 100644 --- a/pkg-r/R/QueryChatGreeter.R +++ b/pkg-r/R/QueryChatGreeter.R @@ -36,19 +36,30 @@ QueryChatGreeter <- R6::R6Class( #' @param data_description Advanced/internal: overrides the QueryChat #' instance's inferred data description for this call only, for the same #' reason as `tables`. + #' + #' For `data_sources` and `data_description`, an explicitly passed value + #' (including `NULL`) is the snapshot and wins over live state; only an + #' omitted argument falls back to live state. build_client = function( base = NULL, tables = NULL, - data_sources = NULL, - data_description = NULL + data_sources, + data_description ) { - private$.client_factory( + args <- list( tables %||% private$.tables, private$.prompt, - base, - data_sources = data_sources, - data_description = data_description + base ) + # Single-bracket assignment: `$<-` would delete the element when the + # snapshot value is NULL, turning an explicit NULL back into "omitted". + if (!missing(data_sources)) { + args["data_sources"] <- list(data_sources) + } + if (!missing(data_description)) { + args["data_description"] <- list(data_description) + } + do.call(private$.client_factory, args) }, #' @description Generate a greeting synchronously and return it as text. diff --git a/pkg-r/tests/testthat/test-server_data_source.R b/pkg-r/tests/testthat/test-server_data_source.R index 946525644..8eee26c8e 100644 --- a/pkg-r/tests/testthat/test-server_data_source.R +++ b/pkg-r/tests/testthat/test-server_data_source.R @@ -241,6 +241,58 @@ describe("QueryChat$server(data_source=) session lifecycle", { }) }) +describe("QueryChat$server(data_source=) retired resource cleanup", { + it("resources replaced while a session is live are cleaned up when the last session ends", { + skip_if_no_dataframe_engine() + local_captured_mod_server() + + qc <- QueryChat$new(NULL, "users", greeting = "Test") + withr::defer(qc$cleanup()) + + session1 <- fake_shiny_session() + session2 <- fake_shiny_session() + qc$server(data_source = new_users_df(), session = session1) + first_source <- qc_data_source(qc, "users") + first_executor <- qc$.__enclos_env__$private$.query_executor + source_cleaned <- spy_on_cleanup(first_source) + executor_cleaned <- spy_on_cleanup(first_executor) + + qc$server(data_source = new_users_df(), session = session2) + + session1$end() # session 2 still live: nothing cleaned yet + expect_false(source_cleaned()) + expect_false(executor_cleaned()) + + session2$end() # last live session: retired resources are flushed + expect_true(source_cleaned()) + expect_true(executor_cleaned()) + }) + + it("retired resources are also cleaned up by $cleanup()", { + skip_if_no_dataframe_engine() + local_captured_mod_server() + + qc <- QueryChat$new(NULL, "users", greeting = "Test") + + qc$server(data_source = new_users_df(), session = fake_shiny_session()) + first_source <- qc_data_source(qc, "users") + source_cleaned <- spy_on_cleanup(first_source) + + qc$server(data_source = new_users_df(), session = fake_shiny_session()) + expect_false(source_cleaned()) + + qc$cleanup() + expect_true(source_cleaned()) + }) + + it("an invalid deferred table_name fails fast at construction", { + expect_error( + QueryChat$new(NULL, table_name = "bad-name"), + "valid SQL table name" + ) + }) +}) + describe("QueryChat$server(data_source=) greeting snapshot", { it("passes a greeting_tables snapshot to mod_server", { skip_if_no_dataframe_engine() @@ -306,6 +358,27 @@ describe("QueryChat$server(data_source=) greeting snapshot", { expect_match(snapshot, "snapshot description", fixed = TRUE) expect_no_match(snapshot, "live description", fixed = TRUE) }) + + it("an explicit NULL data_description snapshot does not fall back to live state", { + skip_if_no_dataframe_engine() + + qc <- QueryChat$new( + new_users_df(), + "users", + greeting = "Test", + data_description = "live description", + client = mock_ellmer_chat_client() + ) + withr::defer(qc$cleanup()) + + # A session whose snapshot had no description must not pick up a + # description inferred by a later session's registration. + prompt <- qc$greeter$build_client( + data_description = NULL + )$get_system_prompt() + + expect_no_match(prompt, "live description", fixed = TRUE) + }) }) describe("Mixing config-time $add_table() with $server(data_source=)", { From 9cde4cdb5494126b851cce830d77304ae87db036 Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 12:34:23 -0500 Subject: [PATCH 7/9] fix(r): harden session lifecycle and cleanup against partial failure Address latest Copilot review on #306: - $cleanup() flushes retired resources via on.exit(), so the fallback cleanup path runs even if current executor/source cleanup throws. - $server() counts a session (and arms its end-of-session flush) only after setup fully succeeds, so a failed call on a live session no longer blocks table mutations or defers replacement cleanup forever. - add_or_replace_table() runs the source-compatibility check inside the cleanup-on-error block, so a failed per-session registration doesn't leak the staged source's connection. - Table replacement skips cleanup/retirement when old and new sources wrap the same DBI connection (two sessions passing the same raw connection each get their own wrapper); the replacement inherits ownership instead of being disconnected out from under. Applies to both $server(data_source=) and $add_table(replace=TRUE). - mod_server() forces greeting_tables/greeting_data_description at entry so the deferred greeting callback sees the session's point-in-time snapshot, not a later session's mutation. - QueryChatGreeter docs: client_factory signature now includes data_description. --- pkg-r/R/DBISource.R | 8 ++ pkg-r/R/QueryChat.R | 50 +++++-- pkg-r/R/QueryChatGreeter.R | 2 +- pkg-r/R/querychat_module.R | 7 + pkg-r/tests/testthat/helper-fixtures.R | 6 +- pkg-r/tests/testthat/test-querychat_module.R | 79 ++++++++++ .../tests/testthat/test-server_data_source.R | 135 ++++++++++++++++++ 7 files changed, 273 insertions(+), 14 deletions(-) diff --git a/pkg-r/R/DBISource.R b/pkg-r/R/DBISource.R index 725c02d8c..572ad89d1 100644 --- a/pkg-r/R/DBISource.R +++ b/pkg-r/R/DBISource.R @@ -212,6 +212,14 @@ DBISource <- R6::R6Class( self$execute_query(NULL) }, + #' @description + #' Get the underlying DBI connection + #' + #' @return The DBI connection this source wraps. + get_connection = function() { + private$conn + }, + #' @description #' Disconnect from the database #' diff --git a/pkg-r/R/QueryChat.R b/pkg-r/R/QueryChat.R index 2fd033d66..aa708853a 100644 --- a/pkg-r/R/QueryChat.R +++ b/pkg-r/R/QueryChat.R @@ -192,7 +192,6 @@ QueryChat <- R6::R6Class( other_sources <- private$.data_sources[ names(private$.data_sources) != table_name ] - check_source_compatibility(other_sources, normalized, table_name) next_sources <- private$.data_sources next_sources[[table_name]] <- normalized @@ -200,9 +199,12 @@ QueryChat <- R6::R6Class( private$auto_fill_data_description(next_sources) tryCatch( { + check_source_compatibility(other_sources, normalized, table_name) private$build_system_prompt(data_sources = next_sources) }, error = function(e) { + # A source normalized here (not user-supplied) owns a fresh + # connection; don't leak it when staging fails. if (!inherits(data_source, "DataSource")) { normalized$cleanup() } @@ -212,7 +214,11 @@ QueryChat <- R6::R6Class( old_source <- private$.data_sources[[table_name]] private$.data_sources <- next_sources - if (!is.null(old_source) && !identical(old_source, normalized)) { + if ( + !is.null(old_source) && + !identical(old_source, normalized) && + !shares_underlying_connection(old_source, normalized) + ) { if (cleanup_replaced) { old_source$cleanup() } else { @@ -648,7 +654,8 @@ QueryChat <- R6::R6Class( old_source <- private$.data_sources[[table_name]] if ( !is.null(old_source) && - !identical(old_source, normalized[[table_name]]) + !identical(old_source, normalized[[table_name]]) && + !shares_underlying_connection(old_source, normalized[[table_name]]) ) { old_source$cleanup() } @@ -1194,14 +1201,6 @@ QueryChat <- R6::R6Class( private$require_initialized("$server") - private$.active_sessions <- private$.active_sessions + 1L - session$onSessionEnded(function() { - private$.active_sessions <- private$.active_sessions - 1L - if (private$.active_sessions == 0) { - private$flush_retired_resources() - } - }) - if (is.null(private$.query_executor)) { private$.query_executor <- build_query_executor(private$.data_sources) } @@ -1247,6 +1246,18 @@ QueryChat <- R6::R6Class( greeting_tables = self$greeter$tables, greeting_data_description = private$.data_description ) + + # Count the session (and arm the end-of-session flush) only after setup + # has fully succeeded: a failed $server() call on a still-live session + # must not block table mutations or defer replacement cleanup forever. + private$.active_sessions <- private$.active_sessions + 1L + session$onSessionEnded(function() { + private$.active_sessions <- private$.active_sessions - 1L + if (private$.active_sessions == 0) { + private$flush_retired_resources() + } + }) + result }, @@ -1268,13 +1279,15 @@ QueryChat <- R6::R6Class( #' #' @return Invisibly returns `NULL`. Resources are cleaned up internally. cleanup = function() { + # The retired-resource flush is the fallback cleanup path; make sure it + # runs even if cleaning the current executor or sources throws. + on.exit(private$flush_retired_resources(), add = TRUE) if (!is.null(private$.query_executor)) { private$.query_executor$cleanup() } for (source in private$.data_sources) { source$cleanup() } - private$flush_retired_resources() invisible(NULL) } ), @@ -1576,6 +1589,19 @@ normalize_data_source <- function(data_source, table_name) { ) } +# Do the two sources wrap the same DBI connection? Each registration of a raw +# connection gets its own DBISource wrapper, so two sources can share one +# connection; the replacement then inherits ownership of it, and cleaning the +# replaced wrapper would disconnect the connection out from under it. +shares_underlying_connection <- function(old_source, new_source) { + if ( + !inherits(old_source, "DBISource") || !inherits(new_source, "DBISource") + ) { + return(FALSE) + } + identical(old_source$get_connection(), new_source$get_connection()) +} + normalize_data_dicts <- function(data_dict) { if (is.null(data_dict)) { return(list()) diff --git a/pkg-r/R/QueryChatGreeter.R b/pkg-r/R/QueryChatGreeter.R index 329c57d73..02f4f0169 100644 --- a/pkg-r/R/QueryChatGreeter.R +++ b/pkg-r/R/QueryChatGreeter.R @@ -13,7 +13,7 @@ QueryChatGreeter <- R6::R6Class( .prompt = NULL ), public = list( - #' @param client_factory function(tables, prompt, base, data_sources) returning a configured greeting client. + #' @param client_factory function(tables, prompt, base, data_sources, data_description) returning a configured greeting client. initialize = function(client_factory) { private$.client_factory <- client_factory private$.tables <- character() diff --git a/pkg-r/R/querychat_module.R b/pkg-r/R/querychat_module.R index 780bcc5ed..79066fbbe 100644 --- a/pkg-r/R/querychat_module.R +++ b/pkg-r/R/querychat_module.R @@ -56,6 +56,13 @@ mod_server <- function( greeting_tables = NULL, greeting_data_description = NULL ) { + # These arrive as lazy promises over live QueryChat state and are read only + # inside the deferred greeting callback below; force them now so the + # greeting reflects the point-in-time snapshot captured when this session's + # $server() call ran, not a later session's mutation. + force(greeting_tables) + force(greeting_data_description) + shiny::moduleServer(id, function(input, output, session) { current_table_val <- shiny::reactiveVal(NULL, label = "current_table") diff --git a/pkg-r/tests/testthat/helper-fixtures.R b/pkg-r/tests/testthat/helper-fixtures.R index 6ee167ee9..3bd318010 100644 --- a/pkg-r/tests/testthat/helper-fixtures.R +++ b/pkg-r/tests/testthat/helper-fixtures.R @@ -69,7 +69,11 @@ local_sqlite_connection <- function( temp_db <- withr::local_tempfile(fileext = ".db", .local_envir = env) conn <- DBI::dbConnect(RSQLite::SQLite(), temp_db) - withr::defer(DBI::dbDisconnect(conn), envir = env) + # QueryChat may legitimately disconnect first (e.g. cleanup-on-replace tests) + withr::defer( + if (DBI::dbIsValid(conn)) DBI::dbDisconnect(conn), + envir = env + ) DBI::dbWriteTable(conn, table_name, data, overwrite = TRUE) diff --git a/pkg-r/tests/testthat/test-querychat_module.R b/pkg-r/tests/testthat/test-querychat_module.R index 464fa80c7..3c4b2ac67 100644 --- a/pkg-r/tests/testthat/test-querychat_module.R +++ b/pkg-r/tests/testthat/test-querychat_module.R @@ -773,6 +773,85 @@ test_that("mod_server() builds the auto-generated greeting from the greeter, not ) }) +test_that("mod_server() forces greeting snapshot args at entry, not at greeting time", { + skip_if_no_dataframe_engine() + + ds <- local_data_frame_source(new_test_df()) + executor <- build_query_executor(list(test_table = ds)) + withr::defer(executor$cleanup()) + + client_factory <- function(...) { + structure(list(), class = c("MockChat", "Chat")) + } + + build_client_calls <- list() + fake_greeting_client <- list( + stream_async = function(prompt) "fake-stream" + ) + fake_greeter <- list( + build_client = function( + base = NULL, + tables = NULL, + data_sources = NULL, + data_description = NULL + ) { + build_client_calls[[length(build_client_calls) + 1L]] <<- list( + tables = tables, + data_description = data_description + ) + fake_greeting_client + } + ) + + captured_greeting_arg <- NULL + local_mocked_bindings( + chat_server = function(id, client, greeting = NULL, ...) { + captured_greeting_arg <<- greeting + mock_chat_server_result(client) + }, + chat_greeting = function(content, ...) content, + .package = "shinychat" + ) + local_mock_chat_restore() + + shiny::testServer( + # Indirect through a server function so the snapshot args reach mod_server() + # as lazy promises over these locals, as they do from QueryChat$server() + function(input, output, session) { + snapshot_tables <- "test_table" + snapshot_description <- "original description" + + mod_server( + "inner", + data_sources = list(test_table = ds), + executor = executor, + greeting = NULL, + client = client_factory, + tools = "query", + history = TRUE, + greeter = fake_greeter, + greeting_base = "base-client", + greeting_tables = snapshot_tables, + greeting_data_description = snapshot_description + ) + + # A later session mutates the live state the promises point at, after + # this session's setup but before its greeting is built + snapshot_tables <- c("test_table", "other_table") + snapshot_description <- "mutated description" + }, + { + suppressWarnings(captured_greeting_arg()) + + expect_identical(build_client_calls[[1]]$tables, "test_table") + expect_identical( + build_client_calls[[1]]$data_description, + "original description" + ) + } + ) +}) + test_that("mod_server() chat_update input updates table state", { skip_if_no_dataframe_engine() diff --git a/pkg-r/tests/testthat/test-server_data_source.R b/pkg-r/tests/testthat/test-server_data_source.R index 8eee26c8e..4295be31d 100644 --- a/pkg-r/tests/testthat/test-server_data_source.R +++ b/pkg-r/tests/testthat/test-server_data_source.R @@ -291,6 +291,141 @@ describe("QueryChat$server(data_source=) retired resource cleanup", { "valid SQL table name" ) }) + + it("$cleanup() flushes retired resources even when another cleanup fails", { + skip_if_no_dataframe_engine() + local_captured_mod_server() + + qc <- QueryChat$new(NULL, "users", greeting = "Test") + withr::defer(qc$cleanup()) + + qc$server(data_source = new_users_df(), session = fake_shiny_session()) + first_source <- qc_data_source(qc, "users") + first_cleaned <- spy_on_cleanup(first_source) + + qc$server(data_source = new_users_df(), session = fake_shiny_session()) + + # Make the current source's cleanup fail (once, so teardown can retry) + current_source <- qc_data_source(qc, "users") + unlockBinding("cleanup", current_source) + fail_once <- TRUE + current_source$cleanup <- function() { + if (fail_once) { + fail_once <<- FALSE + stop("boom") + } + invisible(NULL) + } + + expect_error(qc$cleanup(), "boom") + expect_true(first_cleaned()) + }) +}) + +describe("QueryChat$server(data_source=) registration failures", { + it("a failed $server() call does not count as a live session", { + skip_if_no_dataframe_engine() + testthat::local_mocked_bindings( + mod_server = function(...) stop("mod_server failed"), + .package = "querychat" + ) + + qc <- QueryChat$new(NULL, "users", greeting = "Test") + withr::defer(qc$cleanup()) + + expect_error( + qc$server(data_source = new_users_df(), session = fake_shiny_session()), + "mod_server failed" + ) + + # The failed session must not linger in the live-session count, which + # would block config-time mutations forever + expect_no_error(qc$add_table(new_users_df(), "other")) + }) + + it("a failed per-session registration cleans up the staged source", { + skip_if_no_dataframe_engine() + skip_if_not_installed("RSQLite") + local_captured_mod_server() + + qc <- QueryChat$new(new_users_df(), "users", greeting = "Test") + withr::defer(qc$cleanup()) + + # A DBI source can't be added alongside a data-frame source; registration + # fails after the DBISource wrapper has already been staged + db <- local_sqlite_connection(new_users_df(), "dbtable") + expect_error( + qc$server( + data_source = db$conn, + table_name = "dbtable", + session = fake_shiny_session() + ), + "all tables must be the same type" + ) + + # The staged wrapper owned the connection, so failure must not leak it + expect_false(DBI::dbIsValid(db$conn)) + }) +}) + +describe("QueryChat table replacement with a shared DBI connection", { + it("a session replacing a table with the same connection does not retire it", { + skip_if_not_installed("RSQLite") + local_captured_mod_server() + + db <- local_sqlite_connection(new_users_df(), "users") + con <- db$conn + + qc <- QueryChat$new(NULL, "users", greeting = "Test") + withr::defer(qc$cleanup()) + + session1 <- fake_shiny_session() + session2 <- fake_shiny_session() + qc$server(data_source = con, session = session1) + qc$server(data_source = con, session = session2) + + session1$end() + session2$end() + + # The second session's source wraps the same connection, so flushing the + # retired first wrapper must not disconnect it + expect_true(DBI::dbIsValid(con)) + }) + + it("a replaced connection is still disconnected once no session uses it", { + skip_if_not_installed("RSQLite") + local_captured_mod_server() + + db1 <- local_sqlite_connection(new_users_df(), "users") + db2 <- local_sqlite_connection(new_users_df(), "users") + + qc <- QueryChat$new(NULL, "users", greeting = "Test") + withr::defer(qc$cleanup()) + + session1 <- fake_shiny_session() + session2 <- fake_shiny_session() + qc$server(data_source = db1$conn, session = session1) + qc$server(data_source = db2$conn, session = session2) + + session1$end() + session2$end() + + expect_false(DBI::dbIsValid(db1$conn)) + expect_true(DBI::dbIsValid(db2$conn)) + }) + + it("$add_table(replace=TRUE) does not disconnect a shared connection", { + skip_if_not_installed("RSQLite") + + db <- local_sqlite_connection(new_users_df(), "users") + + qc <- QueryChat$new(db$conn, "users", greeting = "Test") + withr::defer(qc$cleanup()) + + qc$add_table(db$conn, "users", replace = TRUE) + + expect_true(DBI::dbIsValid(db$conn)) + }) }) describe("QueryChat$server(data_source=) greeting snapshot", { From 70b126b1d9769b0fc05403466384b9a938c31144 Mon Sep 17 00:00:00 2001 From: cpsievert Date: Sat, 12 Sep 2026 17:36:17 +0000 Subject: [PATCH 8/9] `devtools::document()` (GitHub Actions) --- pkg-r/man/DBISource.Rd | 16 ++++++++++++++++ pkg-r/man/DataFrameSource.Rd | 1 + pkg-r/man/PinSource.Rd | 1 + pkg-r/man/TblSqlSource.Rd | 1 + 4 files changed, 19 insertions(+) diff --git a/pkg-r/man/DBISource.Rd b/pkg-r/man/DBISource.Rd index 539af25a8..f6132774b 100644 --- a/pkg-r/man/DBISource.Rd +++ b/pkg-r/man/DBISource.Rd @@ -42,6 +42,7 @@ db_source$cleanup() \item \href{#method-DBISource-execute_query}{\code{DBISource$execute_query()}} \item \href{#method-DBISource-test_query}{\code{DBISource$test_query()}} \item \href{#method-DBISource-get_data}{\code{DBISource$get_data()}} + \item \href{#method-DBISource-get_connection}{\code{DBISource$get_connection()}} \item \href{#method-DBISource-cleanup}{\code{DBISource$cleanup()}} \item \href{#method-DBISource-clone}{\code{DBISource$clone()}} } @@ -200,6 +201,21 @@ all original table columns (default: \code{FALSE})} } } +\if{html}{\out{
}} +\if{html}{\out{}} +\if{latex}{\out{\hypertarget{method-DBISource-get_connection}{}}} +\subsection{\code{DBISource$get_connection()}}{ + Get the underlying DBI connection + \subsection{Usage}{ + \if{html}{\out{
}} + \preformatted{DBISource$get_connection()} + \if{html}{\out{
}} + } + \subsection{Returns}{ + The DBI connection this source wraps. + } +} + \if{html}{\out{
}} \if{html}{\out{}} \if{latex}{\out{\hypertarget{method-DBISource-cleanup}{}}} diff --git a/pkg-r/man/DataFrameSource.Rd b/pkg-r/man/DataFrameSource.Rd index b7402f4fa..7ee3c9f6f 100644 --- a/pkg-r/man/DataFrameSource.Rd +++ b/pkg-r/man/DataFrameSource.Rd @@ -53,6 +53,7 @@ df_sqlite$cleanup()
  • DataSource$get_data_description()
  • DBISource$execute_query()
    • +
  • DBISource$get_connection()
  • DBISource$get_data()
  • DBISource$get_db_type()
  • DBISource$get_schema()
    • diff --git a/pkg-r/man/PinSource.Rd b/pkg-r/man/PinSource.Rd index 34eacf4a9..88e877cbd 100644 --- a/pkg-r/man/PinSource.Rd +++ b/pkg-r/man/PinSource.Rd @@ -74,6 +74,7 @@ if (rlang::is_installed(c("pins", "duckdb"))) {
    • DBISource$cleanup()
    • DBISource$execute_query()
      • +
    • DBISource$get_connection()
    • DBISource$get_data()
    • DBISource$get_db_type()
    • DBISource$get_schema()
      • diff --git a/pkg-r/man/TblSqlSource.Rd b/pkg-r/man/TblSqlSource.Rd index f15a1cccc..a359a5456 100644 --- a/pkg-r/man/TblSqlSource.Rd +++ b/pkg-r/man/TblSqlSource.Rd @@ -58,6 +58,7 @@ mtcars_source$cleanup() \if{html}{\out{
      Inherited methods
      }} From 3c6332d957b287d68314daefbd526cf9ff408811 Mon Sep 17 00:00:00 2001 From: Carson Date: Sat, 12 Sep 2026 13:14:36 -0500 Subject: [PATCH 9/9] fix(r): make registration atomic and retired-resource flush retryable Address latest Copilot review on #306: - add_or_replace_table() validates source compatibility before mutating .data_description, and restores both description fields on any staging error, so a failed per-session registration leaves the existing registration fully intact. - flush_retired_resources() retains entries whose cleanup fails, so a transient error no longer permanently drops a resource; the next flush (or $cleanup()) retries it. - flush_retired_resources() skips retired sources that share their underlying DBI connection with a current data source: registrations alternating connections (C1 -> C2 -> C1) while sessions are live no longer disconnect the live connection. --- pkg-r/R/QueryChat.R | 32 ++++- .../tests/testthat/test-server_data_source.R | 114 ++++++++++++++++++ 2 files changed, 143 insertions(+), 3 deletions(-) diff --git a/pkg-r/R/QueryChat.R b/pkg-r/R/QueryChat.R index aa708853a..8660e2e9f 100644 --- a/pkg-r/R/QueryChat.R +++ b/pkg-r/R/QueryChat.R @@ -196,13 +196,19 @@ QueryChat <- R6::R6Class( next_sources <- private$.data_sources next_sources[[table_name]] <- normalized - private$auto_fill_data_description(next_sources) + # Snapshot so a staging failure can be rolled back: a failed + # registration must leave the existing registration fully intact. + prev_description <- private$.data_description + prev_description_mode <- private$.data_description_mode tryCatch( { check_source_compatibility(other_sources, normalized, table_name) + private$auto_fill_data_description(next_sources) private$build_system_prompt(data_sources = next_sources) }, error = function(e) { + private$.data_description <- prev_description + private$.data_description_mode <- prev_description_mode # A source normalized here (not user-supplied) owns a fresh # connection; don't leak it when staging fails. if (!inherits(data_source, "DataSource")) { @@ -259,8 +265,28 @@ QueryChat <- R6::R6Class( retired <- private$.retired_resources private$.retired_resources <- list() for (resource in retired) { - # Best-effort: one failing cleanup must not leave the rest open. - tryCatch(resource$cleanup(), error = function(e) NULL) + # Registrations can alternate connections (C1 -> C2 -> C1) while + # sessions are live, so a retired wrapper may share its connection + # with the now-current source, which owns it. Never clean those. + in_use <- any(vapply( + private$.data_sources, + function(source) shares_underlying_connection(resource, source), + logical(1) + )) + if (in_use) { + next + } + # Best-effort: one failing cleanup must not leave the rest open, but + # retain the failed entry so a later flush (or $cleanup()) can retry. + tryCatch( + resource$cleanup(), + error = function(e) { + private$.retired_resources <- c( + private$.retired_resources, + list(resource) + ) + } + ) } }, diff --git a/pkg-r/tests/testthat/test-server_data_source.R b/pkg-r/tests/testthat/test-server_data_source.R index 4295be31d..8a7ea1864 100644 --- a/pkg-r/tests/testthat/test-server_data_source.R +++ b/pkg-r/tests/testthat/test-server_data_source.R @@ -320,6 +320,42 @@ describe("QueryChat$server(data_source=) retired resource cleanup", { expect_error(qc$cleanup(), "boom") expect_true(first_cleaned()) }) + + it("a failed retired-resource cleanup is retained and retried", { + skip_if_no_dataframe_engine() + local_captured_mod_server() + + qc <- QueryChat$new(NULL, "users", greeting = "Test") + withr::defer(qc$cleanup()) + + session1 <- fake_shiny_session() + qc$server(data_source = new_users_df(), session = session1) + first_source <- qc_data_source(qc, "users") + + # Fail the first cleanup attempt (transiently), succeed on retry + unlockBinding("cleanup", first_source) + attempts <- 0 + first_source$cleanup <- function() { + attempts <<- attempts + 1 + if (attempts == 1) { + stop("transient failure") + } + invisible(NULL) + } + + session2 <- fake_shiny_session() + qc$server(data_source = new_users_df(), session = session2) + + session1$end() # session 2 still live: no flush yet + expect_equal(attempts, 0) + + session2$end() # last live session: flush runs, cleanup fails transiently + expect_equal(attempts, 1) + + # The failed resource is retained, so $cleanup()'s flush retries it + qc$cleanup() + expect_equal(attempts, 2) + }) }) describe("QueryChat$server(data_source=) registration failures", { @@ -366,6 +402,55 @@ describe("QueryChat$server(data_source=) registration failures", { # The staged wrapper owned the connection, so failure must not leak it expect_false(DBI::dbIsValid(db$conn)) }) + + it("a failed registration restores the inferred data description", { + skip_if_no_dataframe_engine() + local_captured_mod_server() + + # A data-frame source carrying a description, so registration infers one + described_df_source <- function(description) { + klass <- R6::R6Class( + "DescribedDataFrameSource", + inherit = DataFrameSource, + public = list( + get_data_description = function() description + ) + ) + klass$new(new_users_df(), "users") + } + + qc <- QueryChat$new(NULL, "users", greeting = "Test") + withr::defer(qc$cleanup()) + qc$add_table(described_df_source("original description"), "users") + expect_identical( + qc$.__enclos_env__$private$.data_description, + "original description" + ) + + testthat::local_mocked_bindings( + QueryChatSystemPrompt = list( + new = function(...) stop("prompt build failed") + ), + .package = "querychat" + ) + + replacement <- described_df_source("replacement description") + withr::defer(replacement$cleanup()) + expect_error( + qc$server(data_source = replacement, session = fake_shiny_session()), + "prompt build failed" + ) + + # The failed registration must not disturb the existing description state + expect_identical( + qc$.__enclos_env__$private$.data_description, + "original description" + ) + expect_identical( + qc$.__enclos_env__$private$.data_description_mode, + "inferred" + ) + }) }) describe("QueryChat table replacement with a shared DBI connection", { @@ -414,6 +499,35 @@ describe("QueryChat table replacement with a shared DBI connection", { expect_true(DBI::dbIsValid(db2$conn)) }) + it("registrations alternating between connections keep the live connection open", { + skip_if_not_installed("RSQLite") + local_captured_mod_server() + + db1 <- local_sqlite_connection(new_users_df(), "users") + db2 <- local_sqlite_connection(new_users_df(), "users") + + qc <- QueryChat$new(NULL, "users", greeting = "Test") + withr::defer(qc$cleanup()) + + session1 <- fake_shiny_session() + session2 <- fake_shiny_session() + session3 <- fake_shiny_session() + qc$server(data_source = db1$conn, session = session1) + qc$server(data_source = db2$conn, session = session2) + # Back to db1's connection: the retired first wrapper shares it with the + # now-current source + qc$server(data_source = db1$conn, session = session3) + + session1$end() + session2$end() + session3$end() + + # Flushing retired wrappers must not disconnect the connection the + # current source uses, but db2's retired wrapper is still cleaned up + expect_true(DBI::dbIsValid(db1$conn)) + expect_false(DBI::dbIsValid(db2$conn)) + }) + it("$add_table(replace=TRUE) does not disconnect a shared connection", { skip_if_not_installed("RSQLite")