feat(rag): add DuckDB provider behind the RAG driver abstraction

Phase 3 of the RAG driver abstraction. Adds a `DuckDbProvider` that keeps
vectors and document content in a `.duckdb` sidecar next to the existing
YAML metadata, selected by the `driver: duckdb` field.

- `src/rag/providers/duckdb.rs` (new): vector search via the vss extension
  and keyword search via fts, an all-or-nothing hydration path (a partial
  read is an error, never a shorter map), and an anti-wipe guard that
  refuses the destructive `CREATE OR REPLACE TABLE` when `data.vectors` is
  empty while `data.files` is not and the store still holds rows.
- `src/rag/mod.rs`: `sync_documents` now refreshes `bm25`/`node_to_docs`
  BEFORE the fallible `provider.rebuild_indexes`. `self.data` is already
  mutated by that point, so propagating a provider error afterwards would
  leave the derived in-memory state describing the previous corpus while
  `data` describes the new one. Both rebuilds are pure functions of
  `self.data` and cannot fail, so running them first is always safe.
- `src/config/paths.rs`: sidecar path helpers.
- `src/rag/providers/mod.rs`, `src/config/agent.rs`: driver dispatch and
  RAG cache keying.

Also keeps `RequestContext::rag_key` in lockstep with `rag` at the two
sites that were still missing it, so that a cache insert and its matching
invalidate are structurally incapable of disagreeing:

- `use_agent` assigned `self.rag` from the agent but never set `rag_key`.
  This one was live. Agent RAGs are inserted under `RagKey::Agent(<name>)`,
  so with `rag_key == None` the invalidation guards in `rebuild_rag` and
  `edit_rag_docs` matched nothing and `.rebuild rag` left the stale cache
  entry in place. Worse, a preceding `.rag <name>` left a stale
  `Named(<name>)` key attached to the agent's RAG, pointing the
  invalidation at an unrelated RAG's cache entry. Now mirrors the insert
  key exactly, yielding `None` when the agent has no RAG.
- `exit_agent` cleared `self.rag` but left `rag_key` behind. Latent rather
  than live, since `rebuild_rag`/`edit_rag_docs` both bail on
  `rag.is_none()` before reaching the invalidate guards, but the guards
  that make it unobservable are not the kind of thing to depend on.

Covered by `use_agent_does_not_carry_stale_rag_key`, and by a new
assertion in `exit_agent_clears_all_agent_state`.
This commit is contained in:
2026-08-10 12:51:56 -06:00
parent d734276927
commit 98d3ba4a83
6 changed files with 1364 additions and 30 deletions
+13 -1
View File
@@ -171,7 +171,15 @@ impl Agent {
let rag = app_state
.rag_cache
.load_with(key, || async move {
Rag::init(&app_clone, "rag", &rag_path_clone, &document_paths, abort).await
Rag::init(
&app_clone,
"rag",
&rag_path_clone,
&document_paths,
abort,
false,
)
.await
})
.await?;
Some(rag)
@@ -983,6 +991,10 @@ async fn init_graph_rags(
extractor_model: rag_node.extractor_model.clone(),
extractor_prompt: rag_node.extractor_prompt.clone(),
graph_hops: rag_node.graph_hops,
// Graph-node RAGs are yaml-only: `RagNode` has no `driver` field, so
// there is nothing to forward. The rest-pattern also keeps this literal
// from breaking on future `RagInitConfig` additions.
..Default::default()
};
let fully_specified = config.embedding_model.is_some()
&& config.chunk_size.is_some()
+115 -1
View File
@@ -16,7 +16,7 @@ use anyhow::{Context, Result, anyhow, bail};
use log::LevelFilter;
use std::collections::HashSet;
use std::env;
use std::fs::{read_dir, read_to_string};
use std::fs::{read_dir, read_to_string, remove_file};
use std::path::{Path, PathBuf};
pub fn config_dir() -> PathBuf {
@@ -414,6 +414,11 @@ pub fn list_rags() -> Vec<String> {
for entry in rd.flatten() {
let name = entry.file_name();
if let Some(name) = name.to_string_lossy().strip_suffix(".yaml") {
// Sidecars are not RAGs. `.duckdb` files are already excluded by
// the `.yaml` suffix check above; this rejects `<name>.sbx-mixin`.
if is_rag_sidecar_name(name) {
continue;
}
names.push(name.to_string());
}
}
@@ -424,6 +429,43 @@ pub fn list_rags() -> Vec<String> {
}
}
/// True for the sidecar YAML files that must never be listed or deleted as RAGs.
/// `name` is the already-stripped stem (i.e. after `strip_suffix(".yaml")`).
/// Uses `ends_with`, not `contains('.')`, so a RAG legitimately named "v2.docs" is
/// not rejected.
pub(crate) fn is_rag_sidecar_name(name: &str) -> bool {
name.ends_with(".sbx-mixin")
}
/// Remove every sidecar belonging to RAG `name` in `dir`. Missing files are NOT an
/// error. A failure to remove an EXISTING mixin IS an error and must propagate — a
/// silently-orphaned mixin keeps a sandbox network permission alive after the user
/// believes it is gone. The `.duckdb` orphan is only wasted disk, so its removal
/// failure is ignorable; the asymmetry is deliberate.
///
/// Callers must run this BEFORE unlinking the primary `.yaml`. If the YAML goes first
/// and this then fails, the RAG disappears from `list_rags()` — so the user can no
/// longer select it to retry — while its `allowedDomains` entry keeps being injected
/// into every sandbox launch.
pub(crate) fn remove_rag_sidecars(dir: &Path, name: &str) -> Result<()> {
let duckdb_path = dir.join(format!("{name}.duckdb"));
if duckdb_path.exists() {
let _ = remove_file(&duckdb_path);
}
let mixin_path = dir.join(format!("{name}.sbx-mixin.yaml"));
if mixin_path.exists() {
remove_file(&mixin_path).with_context(|| {
format!(
"Failed to remove the sandbox mixin for RAG '{name}' at '{}'. \
The RAG was NOT deleted so you can retry; this host remains \
whitelisted in the sandbox until the file is removed.",
mixin_path.display()
)
})?;
}
Ok(())
}
pub fn list_macros() -> Vec<String> {
list_file_names(macros_dir(), ".yaml")
}
@@ -846,4 +888,76 @@ mod tests {
}
let _ = fs::remove_dir_all(&root);
}
/// Unique temp dir for the sidecar helper tests. These take `dir: &Path` directly,
/// so no env-var mutation and therefore no `#[serial]` is needed.
fn sidecar_temp_dir(label: &str) -> PathBuf {
let unique = time::SystemTime::now()
.duration_since(time::UNIX_EPOCH)
.unwrap()
.as_nanos();
let root = env::temp_dir().join(format!("coyote-{label}-test-{unique}"));
fs::create_dir_all(&root).unwrap();
root
}
#[test]
fn is_rag_sidecar_name_accepts_dotted_rag_names() {
// A RAG legitimately named "v2.docs" must not be mistaken for a sidecar.
assert!(!is_rag_sidecar_name("v2.docs"));
assert!(!is_rag_sidecar_name("myrag"));
assert!(is_rag_sidecar_name("myrag.sbx-mixin"));
assert!(is_rag_sidecar_name("v2.docs.sbx-mixin"));
}
#[test]
fn remove_rag_sidecars_removes_both() {
let root = sidecar_temp_dir("rag-sidecars-both");
let duckdb = root.join("docs.duckdb");
let mixin = root.join("docs.sbx-mixin.yaml");
fs::write(&duckdb, "db").unwrap();
fs::write(&mixin, "mixin").unwrap();
remove_rag_sidecars(&root, "docs").unwrap();
assert!(!duckdb.exists(), "the .duckdb sidecar must be removed");
assert!(
!mixin.exists(),
"the .sbx-mixin.yaml sidecar must be removed"
);
let _ = fs::remove_dir_all(&root);
}
#[test]
fn remove_rag_sidecars_is_ok_when_absent() {
let root = sidecar_temp_dir("rag-sidecars-absent");
assert!(remove_rag_sidecars(&root, "docs").is_ok());
let _ = fs::remove_dir_all(&root);
}
#[test]
fn remove_rag_sidecars_runs_before_yaml_unlink() {
let root = sidecar_temp_dir("rag-sidecars-order");
let yaml = root.join("docs.yaml");
fs::write(&yaml, "rag").unwrap();
// A non-empty DIRECTORY at the mixin path makes remove_file fail, standing in
// for any real removal failure (permissions, a busy mount).
let mixin = root.join("docs.sbx-mixin.yaml");
fs::create_dir_all(&mixin).unwrap();
fs::write(mixin.join("blocker"), "x").unwrap();
let err = remove_rag_sidecars(&root, "docs").unwrap_err();
assert!(
err.to_string()
.contains("Failed to remove the sandbox mixin"),
"got: {err}"
);
// The whole point of removing sidecars first: the RAG is still on disk, still
// listed, and the deletion is retryable.
assert!(
yaml.exists(),
"the .yaml must survive a sidecar-removal failure so the delete is retryable"
);
let _ = fs::remove_dir_all(&root);
}
}
+137 -18
View File
@@ -142,6 +142,12 @@ pub struct RequestContext {
pub role: Option<Role>,
pub session: Option<Session>,
pub rag: Option<Arc<Rag>>,
/// The cache key `self.rag` was actually inserted under, carried rather than
/// reconstructed. Reconstruction was the bug: the invalidation sites do not have
/// the information needed to rebuild the key (agent RAGs are inserted under the
/// AGENT's name but `rag.name()` is the constant "rag"), so insert and invalidate
/// silently disagreed. `None` for the temp RAG, which bypasses the cache entirely.
pub rag_key: Option<RagKey>,
pub agent: Option<Agent>,
pub last_message: Option<LastMessage>,
@@ -176,6 +182,7 @@ impl RequestContext {
role: None,
session: None,
rag: None,
rag_key: None,
agent: None,
last_message: None,
tool_scope: ToolScope::default(),
@@ -229,6 +236,7 @@ impl RequestContext {
role: None,
session: None,
rag: None,
rag_key: None,
agent: None,
last_message: None,
tool_scope: ToolScope {
@@ -277,6 +285,7 @@ impl RequestContext {
role: self.role.clone(),
session: self.session.clone(),
rag: self.rag.clone(),
rag_key: self.rag_key.clone(),
agent: self.agent.clone(),
last_message: self.last_message.clone(),
tool_scope: self.tool_scope.clone(),
@@ -315,6 +324,7 @@ impl RequestContext {
role: None,
session: None,
rag: None,
rag_key: None,
agent: None,
last_message: None,
tool_scope: ToolScope {
@@ -2554,6 +2564,14 @@ impl RequestContext {
match file_ext {
Some(file_ext) => {
if let Some(name) = name.to_string_lossy().strip_suffix(file_ext) {
// Sidecars are not independently deletable assets.
// Guarded on `kind == "rag"` because this scan is shared
// by all six kinds, and `session`/`macro` also use
// `.yaml`. The helper lives in paths.rs beside
// list_rags() so both filters cannot drift apart.
if kind == "rag" && paths::is_rag_sidecar_name(name) {
continue;
}
names.push(name.to_string());
}
}
@@ -2590,6 +2608,13 @@ impl RequestContext {
match file_ext {
Some(ext) => {
let path = dir.join(format!("{name}{ext}"));
// Sidecars FIRST. If this fails, the .yaml is still on disk, the
// RAG is still listed, and the user can retry. Unlinking the .yaml
// first would make the deletion unretryable while leaving an
// orphaned mixin whitelisting a host in every sandbox launch.
if kind == "rag" {
paths::remove_rag_sidecars(&dir, &name)?;
}
remove_file(&path).with_context(|| {
format!("Failed to delete {kind} at '{}'", path.display())
})?;
@@ -3729,6 +3754,14 @@ impl RequestContext {
.then(|| Arc::new(RwLock::new(Supervisor::new(max_concurrent, max_depth))));
self.rag = agent.rag();
// Keep `rag_key` in lockstep with `rag`. Agent RAGs are cached under
// `RagKey::Agent(<agent name>)` (see `Agent::init`), so mirror that key exactly;
// leaving the previous key in place would let `.rebuild rag` invalidate an
// unrelated RAG's cache entry, and leaving it `None` would invalidate nothing.
self.rag_key = self
.rag
.is_some()
.then(|| RagKey::Agent(agent.name().to_string()));
self.agent = Some(agent);
self.supervisor = supervisor;
self.inbox = None;
@@ -3777,6 +3810,11 @@ impl RequestContext {
self.pending_agents_guardrail_count = 0;
self.todo_list = TodoList::default();
self.rag.take();
// Cleared alongside `rag` so the pair never disagrees: an agent RAG is
// cached under `RagKey::Agent(<agent name>)`, and leaving that key behind
// would outlive the RAG it names. Latent rather than live today only
// because `rebuild_rag`/`edit_rag_docs` bail on `rag.is_none()` first.
self.rag_key = None;
self.discontinuous_last_message();
}
Ok(())
@@ -4087,7 +4125,9 @@ impl RequestContext {
let rag_cache = self.rag_cache();
let working_mode = self.working_mode;
let rag: Arc<Rag> = match rag {
// The key is returned alongside the Rag rather than assigned inside the match:
// `rag_cache` borrows `self`, so writing `self.rag_key` there is E0506.
let (rag, rag_key): (Arc<Rag>, Option<RagKey>) = match rag {
None => {
let rag_path = self.rag_file(super::TEMP_RAG_NAME);
if rag_path.exists() {
@@ -4095,14 +4135,28 @@ impl RequestContext {
format!("Failed to cleanup previous '{}' rag", super::TEMP_RAG_NAME)
})?;
}
Arc::new(Rag::init(&app, super::TEMP_RAG_NAME, &rag_path, &[], abort_signal).await?)
// The temp RAG is never inserted into the cache, so it has no key.
(
Arc::new(
Rag::init(
&app,
super::TEMP_RAG_NAME,
&rag_path,
&[],
abort_signal,
false,
)
.await?,
),
None,
)
}
Some(name) => {
let rag_path = self.rag_file(name);
let key = RagKey::Named(name.to_string());
rag_cache
.load_with(key, || {
let loaded = rag_cache
.load_with(key.clone(), || {
let app = app.clone();
let rag_path = rag_path.clone();
let abort_signal = abort_signal.clone();
@@ -4111,16 +4165,19 @@ impl RequestContext {
if working_mode.is_cmd() {
bail!("Unknown RAG '{name}'");
}
Rag::init(&app, name, &rag_path, &[], abort_signal.clone()).await
Rag::init(&app, name, &rag_path, &[], abort_signal.clone(), true)
.await
} else {
Rag::load(&app, name, &rag_path)
}
}
})
.await?
.await?;
(loaded, Some(key))
}
};
self.rag = Some(rag);
self.rag_key = rag_key;
Ok(())
}
@@ -4161,12 +4218,9 @@ impl RequestContext {
bail!("No changes")
}
let key = if self.agent.is_some() {
RagKey::Agent(rag.name().to_string())
} else {
RagKey::Named(rag.name().to_string())
};
self.rag_cache().invalidate(&key);
if let Some(key) = self.rag_key.clone() {
self.rag_cache().invalidate(&key);
}
rag.refresh_document_paths(
&new_document_paths,
@@ -4194,12 +4248,9 @@ impl RequestContext {
);
}
let key = if self.agent.is_some() {
RagKey::Agent(rag.name().to_string())
} else {
RagKey::Named(rag.name().to_string())
};
self.rag_cache().invalidate(&key);
if let Some(key) = self.rag_key.clone() {
self.rag_cache().invalidate(&key);
}
let document_paths = rag.document_paths().to_vec();
println!(
@@ -4613,6 +4664,49 @@ mod tests {
assert!(ctx.agent.is_none());
assert!(ctx.rag.is_none());
assert_eq!(ctx.rag_key, None);
}
#[test]
#[serial]
fn use_agent_does_not_carry_stale_rag_key() {
let _guard = TestConfigDirGuard::new();
let mut ctx = create_test_ctx();
let app = ctx.app.config.clone();
let agent_name = format!(
"test_agent_{}",
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos()
);
let agent_dir = paths::agent_data_dir(&agent_name);
create_dir_all(&agent_dir).unwrap();
write(
agent_dir.join("config.yaml"),
format!("name: {agent_name}\ninstructions: hi\n"),
)
.unwrap();
// Stand in for the state `.rag docs` leaves behind: `use_rag` sets `rag` and
// `rag_key` together, so a named key is live when the agent is entered.
ctx.rag_key = Some(RagKey::Named("docs".to_string()));
tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap()
.block_on(async {
ctx.use_agent(&app, &agent_name, None, utils::create_abort_signal())
.await
.unwrap();
});
// This agent has no RAG, so `rag` is None and `rag_key` must be None as well.
// Carrying `Named("docs")` across the transition would point `.rebuild rag`
// at an unrelated RAG's cache entry.
assert!(ctx.rag.is_none());
assert_eq!(ctx.rag_key, None);
}
#[test]
@@ -6000,6 +6094,31 @@ mod tests {
assert!(paths::list_rags().is_empty());
}
/// A `<name>.sbx-mixin.yaml` sidecar must not appear as a phantom RAG in TAB
/// completion or `.list rag`. A RAG whose name legitimately contains a dot must
/// still be listed — the filter uses `ends_with`, not `contains('.')`.
#[test]
#[serial]
fn list_rags_skips_sbx_mixin_sidecars() {
let _guard = TestConfigDirGuard::new();
let rags_dir = paths::rags_dir();
create_dir_all(&rags_dir).unwrap();
write(rags_dir.join("docs.yaml"), "embedding_model: test").unwrap();
write(rags_dir.join("docs.sbx-mixin.yaml"), "kind: mixin").unwrap();
write(rags_dir.join("v2.docs.yaml"), "embedding_model: test").unwrap();
let names = paths::list_rags();
assert!(names.contains(&"docs".to_string()));
assert!(
names.contains(&"v2.docs".to_string()),
"a dotted RAG name must still be listed: {names:?}"
);
assert!(
!names.contains(&"docs.sbx-mixin".to_string()),
"the sandbox mixin sidecar must not appear as a RAG: {names:?}"
);
}
#[test]
#[serial]
fn use_agent_errors_when_already_in_session() {