From 35d7578842c93d49c8cbdafb68dc27ce4540f353 Mon Sep 17 00:00:00 2001 From: Uzayer Masud <16221196+uzayer@users.noreply.github.com> Date: Sat, 10 Oct 2026 00:00:17 +0600 Subject: [PATCH 1/3] fix(codeindex): index the engine, show built-in skips, say when find_symbol has no match, type method receivers - .atlasignore re-includes vendor/atlas-engine (first-party per ADR-0003); a re-include now lifts only the skip it names, so target/ and node_modules inside it stay pruned. The CI planner classifies it as an input of atlas-codeindex's tests. - files pruned by built-in rules (vendor, .atlasignore) count as skipped - find_symbol leads with 'no exact match for "X"' when nothing matches the name - the resolver infers Rust receiver types (let bindings, tuples, fields, method chains, Arc/Box/guards, ?/await) and links a method only to that type's impl; a name-only match in another crate drops to 0.375 confidence - eval_edges reports call-only recall beside the historical metric; a gold answer gains the moved memory store - examples get Metal on macOS like the app does Call recall vs rust-analyzer: 0.849 -> 0.926, precision 0.972 -> 0.980. On the PR #382 metric (every reference on a call line, types included) 0.661 -> 0.700; that metric is capped near 0.70. EXTRACTOR_VERSION 2.3: indexes rebuild once. Co-authored-by: Ahammad Nafiz --- .atlasignore | 7 + crates/atlas-codeindex/Cargo.toml | 8 + crates/atlas-codeindex/eval/questions.jsonl | 2 +- .../atlas-codeindex/examples/build_index.rs | 10 +- crates/atlas-codeindex/examples/eval_edges.rs | 164 +++++- .../examples/eval_retrieval.rs | 2 + crates/atlas-codeindex/src/graph_batch.rs | 84 ++-- crates/atlas-codeindex/src/graph_extract.rs | 257 +++++++++- crates/atlas-codeindex/src/import_resolve.rs | 29 +- crates/atlas-codeindex/src/lib.rs | 3 +- crates/atlas-codeindex/src/query.rs | 26 +- crates/atlas-codeindex/src/resolve.rs | 468 +++++++++++++++++- crates/atlas-codeindex/src/rust_crates.rs | 41 +- crates/atlas-codeindex/src/rust_types.rs | 305 ++++++++++++ crates/atlas-codeindex/src/scan.rs | 193 ++++++-- crates/atlas-codeindex/src/skip.rs | 114 ++++- crates/atlas-codeindex/src/store.rs | 2 +- crates/atlas-codeindex/src/tests/build.rs | 7 +- crates/atlas-codeindex/src/tests/graph.rs | 243 +++++++++ crates/atlas-codeindex/src/tests/query.rs | 25 + crates/atlas-codeindex/src/tests/scan.rs | 51 +- crates/atlas-codeindex/src/tests/skip.rs | 57 +++ crates/atlas-codeindex/src/universe.rs | 26 +- scripts/ci-affected.mjs | 8 + src-tauri/src/commands/code_index/symbols.rs | 36 +- 25 files changed, 2000 insertions(+), 168 deletions(-) create mode 100644 .atlasignore create mode 100644 crates/atlas-codeindex/src/rust_types.rs diff --git a/.atlasignore b/.atlasignore new file mode 100644 index 000000000..8b2dbdbbe --- /dev/null +++ b/.atlasignore @@ -0,0 +1,7 @@ +# What Atlas's code index (symbols, chunks, embeddings) leaves out beyond +# .gitignore and its built-in skips. Gitignore syntax; a `!` line re-includes +# something a built-in rule skipped. + +# The engine is first-party code (ADR-0003), not a vendored dependency: index +# it despite the built-in `vendor/` skip. +!vendor/atlas-engine/ diff --git a/crates/atlas-codeindex/Cargo.toml b/crates/atlas-codeindex/Cargo.toml index f5b75754a..68650a168 100644 --- a/crates/atlas-codeindex/Cargo.toml +++ b/crates/atlas-codeindex/Cargo.toml @@ -56,6 +56,14 @@ atlas-embed = { path = "../atlas-embed" } scip = "0.10" protobuf = "3" +# On macOS the examples embed on the Apple-Silicon GPU, as the app does +# (`src-tauri/Cargo.toml` turns the same feature on for the same target): +# on the CPU, eval_retrieval embeds about 2 chunks a second. Additive over the +# plain dev-dependency above; Linux and Windows builds never see it, and +# `gpu_device()` falls back to the CPU when Metal fails to start. +[target.'cfg(target_os = "macos")'.dev-dependencies] +atlas-embed = { path = "../atlas-embed", features = ["metal"] } + # The workspace lint table is only real for members that opt in — # without this stanza it is dead configuration (#60). [lints] diff --git a/crates/atlas-codeindex/eval/questions.jsonl b/crates/atlas-codeindex/eval/questions.jsonl index 6416d7d41..7058473df 100644 --- a/crates/atlas-codeindex/eval/questions.jsonl +++ b/crates/atlas-codeindex/eval/questions.jsonl @@ -25,6 +25,6 @@ {"q":"send a UI action to the window and wait for its reply with a timeout","gold":["src-tauri/src/commands/ui_server/bridge.rs"]} {"q":"app configuration defaults and per-project overrides","gold":["src-tauri/src/state/atlas_config.rs"]} {"q":"map memory document ids to vector keys with content hashes","gold":["crates/atlas-memory/src/manifest.rs"]} -{"q":"shared memory sqlite store per git scope","gold":["src-tauri/src/commands/shared_memory.rs"]} +{"q":"shared memory sqlite store per git scope","gold":["src-tauri/src/commands/shared_memory.rs","crates/atlas-memory/src/record.rs"]} {"q":"spawn a pseudo terminal for a shell","gold":["crates/atlas-terminal/src/lib.rs"]} {"q":"time out memory retrieval so a tool call never blocks","gold":["src-tauri/src/commands/memory_retrieve.rs"]} diff --git a/crates/atlas-codeindex/examples/build_index.rs b/crates/atlas-codeindex/examples/build_index.rs index 7cc889573..fa4db5e44 100644 --- a/crates/atlas-codeindex/examples/build_index.rs +++ b/crates/atlas-codeindex/examples/build_index.rs @@ -34,6 +34,9 @@ fn main() -> Result<(), Box> { for (reason, n) in ix.status()?.skipped { println!(" skipped {reason}: {n}"); } + for (reason, dir, n) in atlas_codeindex::skipped_by_dir(&stats.skipped) { + println!(" {reason} {dir}: {n}"); + } let t = Instant::now(); let st = ix.reconcile(&CancelToken::new())?; @@ -47,17 +50,20 @@ fn main() -> Result<(), Box> { } let t = Instant::now(); - let (hits, page) = ix.find_symbol(&SymbolQuery { + let (hits, page, exact) = ix.find_symbol_with_exact(&SymbolQuery { query: query.clone(), limit: 5, ..SymbolQuery::default() })?; println!( - "find_symbol {query:?}: {} of {} in {:?}", + "find_symbol {query:?}: {} of {} ({exact} exact) in {:?}", hits.len(), page.total, t.elapsed() ); + if exact == 0 && !hits.is_empty() { + println!(" no exact match for {query:?}; closest by words:"); + } for h in hits { println!( " {} {} {}:{}-{}", diff --git a/crates/atlas-codeindex/examples/eval_edges.rs b/crates/atlas-codeindex/examples/eval_edges.rs index 997de09e3..bad43b837 100644 --- a/crates/atlas-codeindex/examples/eval_edges.rs +++ b/crates/atlas-codeindex/examples/eval_edges.rs @@ -1,11 +1,29 @@ //! Compare atlas-codeindex call edges with rust-analyzer's SCIP index. -//! Usage: rust-analyzer scip . && cargo run -p atlas-codeindex --release --example eval_edges -- . index.scip +//! Usage: rust-analyzer scip . && cargo run -p atlas-codeindex --release --example eval_edges -- . index.scip [truth.tsv] +//! +//! Only edges from files the SCIP index covers are judged, so a SCIP index of part of a +//! workspace (vendored crates emptied, say) measures that part. `truth.tsv` (optional) gets +//! every SCIP reference to a project definition: `ref_rel line col symbol def_rel def_line`. +//! +//! Definitions are compared by line: SCIP's is the identifier's, ours starts at the item's +//! first attribute or doc comment, so a target matches at ±3 lines of our start or on the +//! line that names it. SCIP gives every test target of a package one symbol namespace, so +//! two `tests/*.rs` with an `fn helper` share a symbol: a reference prefers the definition +//! in its own file and otherwise accepts any of them. +//! +//! Recall is printed three ways. "at our reference sites" is the historical number: every +//! SCIP reference on a line where we wrote a call edge, types, fields, modules and enum +//! variants included, which a call graph never links. "call recall" counts only references +//! to functions and methods (`().` symbols): at our sites, then at every site SCIP saw, +//! which also counts non-call mentions (`use` lines, `map(f)`). use std::collections::{HashMap, HashSet}; use std::path::Path; use protobuf::Message; use rusqlite::Connection; +type Site = (String, i32); + fn main() { let args: Vec = std::env::args().collect(); let (root, scip_path) = (Path::new(&args[1]), Path::new(&args[2])); @@ -14,12 +32,14 @@ fn main() { .expect("build"); let index = scip::types::Index::parse_from_bytes(&std::fs::read(scip_path).expect("read scip")) .expect("parse scip"); - // SCIP: definition location of every symbol, and every reference occurrence (file, line) -> symbol. - let mut def_at: HashMap = HashMap::new(); - let mut refs: HashMap<(String, i32), HashSet> = HashMap::new(); + // SCIP: definition locations of every symbol, and every reference occurrence. + let mut def_at: HashMap> = HashMap::new(); + let mut refs: Vec<(Site, i32, String)> = Vec::new(); + let mut docs: HashSet = HashSet::new(); for doc in &index.documents { // rust-analyzer on Windows writes `\` separators; the index uses `/`. let rel = doc.relative_path.replace('\\', "/"); + docs.insert(rel.clone()); for occ in &doc.occurrences { // `local N` symbols are numbered per document: never a cross-file definition. if occ.symbol.is_empty() || occ.symbol.starts_with("local ") { @@ -27,49 +47,145 @@ fn main() { } let line = occ.range.first().copied().unwrap_or(0) + 1; if occ.symbol_roles & (scip::types::SymbolRole::Definition as i32) != 0 { - def_at.insert(occ.symbol.clone(), (rel.clone(), line)); - } else { - refs.entry((rel.clone(), line)) + def_at + .entry(occ.symbol.clone()) .or_default() - .insert(occ.symbol.clone()); + .push((rel.clone(), line)); + } else { + let col = occ.range.get(1).copied().unwrap_or(0); + refs.push(((rel.clone(), line), col, occ.symbol.clone())); } } } - let truth: HashSet<((String, i32), (String, i32))> = refs + // One truth group per (site, symbol): the definitions that reference may mean. + let mut groups: HashMap<(Site, String), Vec> = HashMap::new(); + let mut rows: Vec = Vec::new(); + for (at, col, sym) in &refs { + let Some(defs) = def_at.get(sym) else { + continue; + }; + let own: Vec = defs.iter().filter(|d| d.0 == at.0).cloned().collect(); + let defs = if own.is_empty() { defs.clone() } else { own }; + for d in &defs { + rows.push(format!( + "{}\t{}\t{col}\t{sym}\t{}\t{}", + at.0, at.1, d.0, d.1 + )); + } + groups.entry((at.clone(), sym.clone())).or_insert(defs); + } + if let Some(out) = args.get(3) { + rows.sort(); + rows.dedup(); + std::fs::write(out, rows.join("\n") + "\n").expect("write truth"); + } + let truth: HashSet<(Site, Site)> = groups .iter() - .flat_map(|(at, syms)| { - syms.iter() - .filter_map(|s| def_at.get(s).map(|d| (at.clone(), d.clone()))) - }) + .flat_map(|((at, _), defs)| defs.iter().map(|d| (at.clone(), d.clone()))) .collect(); + let db = Connection::open(root.join(".atlas/code-index/index.db")).expect("db"); let mut stmt = db .prepare( - "SELECT rf.rel, r.line, df.rel, d.start_line FROM edges e JOIN refs r ON r.id = e.ref_id - JOIN files rf ON rf.id = r.file_id JOIN symbols d ON d.id = e.dst_symbol_id JOIN files df ON df.id = d.file_id + "SELECT rf.rel, r.line, df.rel, d.start_line, d.end_line, d.name FROM edges e + JOIN refs r ON r.id = e.ref_id JOIN files rf ON rf.id = r.file_id + JOIN symbols d ON d.id = e.dst_symbol_id JOIN files df ON df.id = d.file_id WHERE e.kind = 'call' AND rf.lang = 'rust'", ) .expect("query"); - let ours: Vec<((String, i32), (String, i32))> = stmt - .query_map([], |r| Ok(((r.get(0)?, r.get(1)?), (r.get(2)?, r.get(3)?)))) + let mut lines: HashMap> = HashMap::new(); + // Each edge: its site and the definition lines it may match. + let ours: Vec<(Site, String, Vec)> = stmt + .query_map([], |r| { + Ok(( + (r.get::<_, String>(0)?, r.get::<_, i32>(1)?), + r.get::<_, String>(2)?, + r.get::<_, i32>(3)?, + r.get::<_, i32>(4)?, + r.get::<_, String>(5)?, + )) + }) .expect("rows") .filter_map(Result::ok) + .filter(|o| docs.contains(&o.0 .0)) + .map(|(at, drel, start, end, name)| { + let mut cand: Vec = (start - 3..=start + 3).collect(); + let text = lines.entry(drel.clone()).or_insert_with(|| { + std::fs::read_to_string(root.join(&drel)) + .unwrap_or_default() + .lines() + .map(str::to_string) + .collect() + }); + let named = (start..=end).find(|l| { + text.get(*l as usize - 1).is_some_and(|t| { + let t = t.trim_start(); + !t.starts_with("//") + && !t.starts_with("#[") + && t.split(|c: char| !(c.is_alphanumeric() || c == '_')) + .any(|w| w == name) + }) + }); + cand.extend(named); + (at, drel, cand) + }) .collect(); - // A def line may differ by attributes/doc comments: accept ±3 lines. - let hit = |o: &((String, i32), (String, i32))| { - (-3..=3).any(|d| truth.contains(&(o.0.clone(), (o.1 .0.clone(), o.1 .1 + d)))) + let hit = |o: &(Site, String, Vec)| { + o.2.iter() + .any(|l| truth.contains(&(o.0.clone(), (o.1.clone(), *l)))) }; let tp = ours.iter().filter(|o| hit(o)).count(); - let ref_lines: HashSet<&(String, i32)> = ours.iter().map(|o| &o.0).collect(); - let truth_at_our_ref_sites = truth.iter().filter(|t| ref_lines.contains(&t.0)).count(); println!( "edges: ours {} | matched {} | precision {:.3}", ours.len(), tp, tp as f64 / ours.len().max(1) as f64 ); + + let ref_lines: HashSet<&Site> = ours.iter().map(|o| &o.0).collect(); + let mut found: HashSet<(Site, Site)> = HashSet::new(); + for o in &ours { + for l in &o.2 { + found.insert((o.0.clone(), (o.1.clone(), *l))); + } + } + let satisfied = |at: &Site, defs: &[Site]| { + defs.iter() + .any(|d| found.contains(&(at.clone(), d.clone()))) + }; + let (mut all_n, mut all_hit, mut site_n, mut site_hit) = (0, 0, 0, 0); + let (mut call_n, mut call_hit, mut call_site_n, mut call_site_hit) = (0, 0, 0, 0); + for ((at, sym), defs) in &groups { + let ok = satisfied(at, defs); + let at_ours = ref_lines.contains(at); + let callable = sym.ends_with("()."); + if at_ours { + site_n += 1; + site_hit += ok as usize; + } + all_n += 1; + all_hit += ok as usize; + if callable { + call_n += 1; + call_hit += ok as usize; + if at_ours { + call_site_n += 1; + call_site_hit += ok as usize; + } + } + } + let r = |a: usize, b: usize| a as f64 / b.max(1) as f64; + println!( + "recall at our reference sites (every symbol kind): {:.3} ({site_hit}/{site_n}; all sites {:.3})", + r(site_hit, site_n), + r(all_hit, all_n) + ); + println!( + "call recall at our reference sites: {:.3} ({call_site_hit}/{call_site_n})", + r(call_site_hit, call_site_n) + ); println!( - "recall at our reference sites: {:.3}", - tp as f64 / truth_at_our_ref_sites.max(1) as f64 + "call recall at every SCIP reference site: {:.3} ({call_hit}/{call_n})", + r(call_hit, call_n) ); } diff --git a/crates/atlas-codeindex/examples/eval_retrieval.rs b/crates/atlas-codeindex/examples/eval_retrieval.rs index 163121c48..303e70bfd 100644 --- a/crates/atlas-codeindex/examples/eval_retrieval.rs +++ b/crates/atlas-codeindex/examples/eval_retrieval.rs @@ -3,6 +3,8 @@ //! `model_dir` is a model downloaded through the app. A code model's directory gets the preset //! Atlas loads it with, as the app rewrites it on every load; any other must hold its own //! `atlas-embed.json`. +//! On macOS the model runs on Metal with no extra flags (a target-specific dev-dependency in +//! this crate's Cargo.toml); elsewhere it embeds on the CPU. use std::path::Path; use atlas_codeindex::{CodeIndex, SemanticQuery}; diff --git a/crates/atlas-codeindex/src/graph_batch.rs b/crates/atlas-codeindex/src/graph_batch.rs index 0b539cb45..58f30d81b 100644 --- a/crates/atlas-codeindex/src/graph_batch.rs +++ b/crates/atlas-codeindex/src/graph_batch.rs @@ -17,7 +17,9 @@ use std::path::{Path, PathBuf}; use rusqlite::{params, Connection, OptionalExtension}; use crate::graph_extract::{GraphExtract, RefKind}; -use crate::import_resolve::{compute_modules, resolve_imports, ProjectConfig}; +use crate::import_resolve::{ + compute_modules, install_module_aliases, resolve_imports, ProjectConfig, +}; use crate::importance; use crate::resolve::{RefRow, Resolver}; use crate::rust_crates::CrateGraph; @@ -129,19 +131,7 @@ impl GraphBatch { let Some((fid, module, surface)) = row else { return Ok(()); }; - let mut names: BTreeMap> = BTreeMap::new(); - let mut stmt = - conn.prepare("SELECT name, kind, signature FROM symbols WHERE file_id = ?1")?; - for row in stmt.query_map([fid], |r| { - Ok(( - r.get::<_, String>(0)?, - r.get::<_, String>(1)?, - r.get::<_, String>(2)?, - )) - })? { - let (n, k, s) = row?; - names.entry(n).or_default().insert((k, s)); - } + let names = surface_names(conn, fid)?; let keys = symbol_keys(conn, fid)?; let mut stmt = conn.prepare( "SELECT e.ref_id, e.src_symbol_id, e.kind, e.confidence, e.strategy, e.dst_symbol_id @@ -277,7 +267,7 @@ impl GraphBatch { fn resolve_incremental(self, conn: &Connection, root: &Path) -> Result { let mut u = Universe::load(conn)?; let mut cfg = ProjectConfig::new(root, load_crates(conn, root, false)?); - let modules = compute_modules(&u, &cfg); + let (modules, aliases) = compute_modules(&u, &cfg); let moved = u.files.iter().zip(&modules).any(|(f, m)| { let before = if self.written.contains(&f.rel) { self.pre.get(&f.rel).map(|p| p.module.as_str()) @@ -292,6 +282,7 @@ impl GraphBatch { } write_modules(conn, &u, &modules)?; u.set_modules(modules); + install_module_aliases(&mut u, &aliases); let targets = resolve_imports(&u, &mut cfg); // A rewritten file may come back under a new id; importers' rows were SET NULL by the // cascade and now point at the same file again — not a binding change. @@ -351,19 +342,7 @@ impl GraphBatch { r.get(0) })?; if surface != pre.surface { - let mut now: BTreeMap> = BTreeMap::new(); - let mut stmt = - conn.prepare("SELECT name, kind, signature FROM symbols WHERE file_id = ?1")?; - for row in stmt.query_map([fid], |r| { - Ok(( - r.get::<_, String>(0)?, - r.get::<_, String>(1)?, - r.get::<_, String>(2)?, - )) - })? { - let (n, k, s) = row?; - now.entry(n).or_default().insert((k, s)); - } + let now = surface_names(conn, fid)?; for n in pre.names.keys().chain(now.keys()) { if pre.names.get(n) != now.get(n) { names_delta.insert(n.clone()); @@ -409,6 +388,10 @@ impl GraphBatch { "WHERE id IN (SELECT value FROM json_each(?1))", Some(&json_ids(&redo_refs)), )?); + // A typed chain (`~…`) reads return types and fields of files it never names. + if !names_delta.is_empty() { + refs.extend(load_refs(conn, "WHERE receiver LIKE '~%'", None)?); + } refs.sort_by_key(|r| r.id); refs.dedup_by_key(|r| r.id); @@ -439,9 +422,10 @@ pub(crate) fn resolve_all( ) -> Result { let mut u = Universe::load(conn)?; let mut cfg = ProjectConfig::new(root, load_crates(conn, root, refresh_crates)?); - let modules = compute_modules(&u, &cfg); + let (modules, aliases) = compute_modules(&u, &cfg); write_modules(conn, &u, &modules)?; u.set_modules(modules); + install_module_aliases(&mut u, &aliases); let targets = resolve_imports(&u, &mut cfg); write_import_targets(conn, &u, &targets, &BTreeSet::new())?; conn.execute("DELETE FROM edges", [])?; @@ -599,8 +583,34 @@ fn symbol_keys( Ok(out) } +/// Name → `(kind, signature)` of a file's symbols, plus `("field", "name: type")` under each +/// struct's name: what a change to is a change other files' resolution may see. +fn surface_names( + conn: &Connection, + fid: i64, +) -> Result>, IndexError> { + let mut names: BTreeMap> = BTreeMap::new(); + let mut stmt = conn.prepare( + "SELECT name, kind, signature FROM symbols WHERE file_id = ?1 + UNION ALL + SELECT s.name, 'field', r.name || ': ' || r.receiver FROM refs r + JOIN symbols s ON s.id = r.src_symbol_id WHERE r.file_id = ?1 AND r.kind = 'field'", + )?; + for row in stmt.query_map([fid], |r| { + Ok(( + r.get::<_, String>(0)?, + r.get::<_, String>(1)?, + r.get::<_, String>(2)?, + )) + })? { + let (n, k, s) = row?; + names.entry(n).or_default().insert((k, s)); + } + Ok(names) +} + /// blake3 over what other files' resolution can see: symbols (qn, kind, signature, exported), -/// `pub` imports and `mod` declarations. +/// struct fields, `pub` imports and `mod` declarations. fn surface_hash(conn: &Connection, fid: i64) -> Result, IndexError> { let mut h = blake3::Hasher::new(); let mut stmt = conn.prepare( @@ -617,6 +627,20 @@ fn surface_hash(conn: &Connection, fid: i64) -> Result, IndexError> { let (k, qn, sig, exp) = row?; h.update(format!("s\t{k}\t{qn}\t{sig}\t{exp}\n").as_bytes()); } + let mut stmt = conn.prepare( + "SELECT s.qualified_name, r.name, r.receiver FROM refs r JOIN symbols s ON s.id = r.src_symbol_id + WHERE r.file_id = ?1 AND r.kind = 'field' ORDER BY r.start_byte", + )?; + for row in stmt.query_map([fid], |r| { + Ok(( + r.get::<_, String>(0)?, + r.get::<_, String>(1)?, + r.get::<_, String>(2)?, + )) + })? { + let (owner, n, ty) = row?; + h.update(format!("f\t{owner}\t{n}\t{ty}\n").as_bytes()); + } let mut stmt = conn.prepare( "SELECT local_name, module_path, imported_name FROM imports WHERE file_id = ?1 AND is_pub = 1 ORDER BY local_name, module_path, imported_name", )?; diff --git a/crates/atlas-codeindex/src/graph_extract.rs b/crates/atlas-codeindex/src/graph_extract.rs index f695fe757..a26035cbd 100644 --- a/crates/atlas-codeindex/src/graph_extract.rs +++ b/crates/atlas-codeindex/src/graph_extract.rs @@ -9,6 +9,14 @@ //! Python `self.x()`, a Go method receiver `r.x()` and a local whose type we saw //! (`let v = Foo::new(); v.x()`) are all rewritten to `"::"`. //! - `"x.y"` value receiver: a member call on an expression we could not type. +//! - `"~~"` (Rust) a receiver whose type depends on other files: a local +//! bound from a call, a field, a method chain or a generic type such as `Arc`. +//! `` is `|`-joined steps the resolver evaluates (`resolve::Resolver::chain_type`): +//! `T`, `F#` (a call's result), `.` (a method call's +//! result), `:`, `?`, `^` (`.await`) and `#` (tuple element). `` is the +//! receiver as it would otherwise read, used when the chain cannot be typed. +//! - field refs (Rust): one per struct field, `name` = the field (`0`, `1`, … for a tuple +//! struct), `receiver` = its type text; never an edge, read by receiver typing. //! - inherit/impl refs: `receiver = "\t"`, qualifier encoded as above. //! - Imports: see [`RawImport`]. @@ -23,6 +31,7 @@ pub enum RefKind { Inherit, Impl, Value, + Field, } impl RefKind { @@ -33,6 +42,7 @@ impl RefKind { Self::Inherit => "inherit", Self::Impl => "impl", Self::Value => "value", + Self::Field => "field", } } @@ -43,6 +53,7 @@ impl RefKind { "inherit" => Some(Self::Inherit), "impl" => Some(Self::Impl), "value" => Some(Self::Value), + "field" => Some(Self::Field), _ => None, } } @@ -306,6 +317,10 @@ const DECL_KINDS: &[&str] = &[ "type_alias", ]; const MAX_RECEIVER_BYTES: usize = 128; +/// Longest typed chain kept (`~…~…` receivers, field types); longer ones are not typed. +const MAX_CHAIN_BYTES: usize = 240; +/// Most steps in one typed chain. +const MAX_CHAIN_STEPS: usize = 8; struct Frame { node_id: usize, @@ -450,11 +465,19 @@ impl<'a> Walker<'a> { } if recv.kind() == "identifier" { if let Some(ty) = self.local_type(t) { + if ty.starts_with('~') { + return format!("{ty}~{}", compact(t)); + } if !ty.is_empty() { return format!("{ty}::"); } } } + if self.lang == GLang::Rust && recv.kind() != "identifier" { + if let Some(chain) = self.rust_chain(recv, 0) { + return format!("~{}~{}", chain.join("|"), compact(t)); + } + } compact(t) } @@ -512,6 +535,47 @@ impl<'a> Walker<'a> { } fn bind_pattern(&mut self, pattern: Node, ty: &str) { + // `let (a, b) = f()` with a typed chain: each element is `…|#i`; + // `let Some(x) = f() else` / `let Ok(x) = …`: `x` is the chain peeled with `?`. + if let Some(chain) = ty.strip_prefix('~') { + match pattern.kind() { + "tuple_pattern" => { + let mut c = pattern.walk(); + let kids: Vec = pattern.named_children(&mut c).collect(); + for (i, k) in kids.into_iter().enumerate() { + if let Some(id) = plain_binding(k) { + let name = self.text(id); + self.bind(name, &format!("~{chain}|#{i}")); + } else { + self.bind_pattern(k, ""); + } + } + return; + } + "tuple_struct_pattern" => { + let ctor = pattern + .child_by_field_name("type") + .map(|t| self.text(t)) + .unwrap_or(""); + let mut c = pattern.walk(); + let kids: Vec = pattern + .named_children(&mut c) + .filter(|k| { + Some(k.id()) != pattern.child_by_field_name("type").map(|t| t.id()) + }) + .collect(); + if let (true, [k]) = (matches!(ctor, "Some" | "Ok"), kids.as_slice()) { + let Some(id) = plain_binding(*k) else { + return self.bind_pattern(*k, ""); + }; + let name = self.text(id); + self.bind(name, &format!("~{chain}|?")); + return; + } + } + _ => {} + } + } let mut ids = Vec::new(); pattern_idents(pattern, &mut ids, 0); let single = ids.len() == 1; @@ -606,11 +670,10 @@ impl<'a> Walker<'a> { "parameter" => { let ty = p .child_by_field_name("type") - .and_then(rust_type_leaf) - .map(|l| self.text(l)) - .unwrap_or(""); + .map(|t| self.rust_declared_type(t)) + .unwrap_or_default(); if let Some(pat) = p.child_by_field_name("pattern") { - self.bind_pattern(pat, ty); + self.bind_pattern(pat, &ty); } } "identifier" => { @@ -625,14 +688,45 @@ impl<'a> Walker<'a> { } "let_declaration" => { let ty = if let Some(t) = n.child_by_field_name("type") { - rust_type_leaf(t).map(|l| self.text(l)).unwrap_or("") + self.rust_declared_type(t) } else { n.child_by_field_name("value") - .map(|v| self.rust_value_type(v)) - .unwrap_or("") + .map(|v| { + let plain = self.rust_value_type(v); + if !plain.is_empty() { + return plain.to_string(); + } + self.rust_chain(v, 0) + .map(|c| format!("~{}", c.join("|"))) + .unwrap_or_default() + }) + .unwrap_or_default() }; if let Some(pat) = n.child_by_field_name("pattern") { - self.bind_pattern(pat, ty); + self.bind_pattern(pat, &ty); + } + true + } + "field_declaration" if in_struct(n) => { + if let (Some(name), Some(ty)) = + (n.child_by_field_name("name"), n.child_by_field_name("type")) + { + if let Some(t) = chain_text(self.text(ty)) { + let field = self.text(name); + self.emit(RefKind::Field, name, field, t); + } + } + true + } + "ordered_field_declaration_list" + if n.parent().is_some_and(|p| p.kind() == "struct_item") => + { + let mut c = n.walk(); + let types: Vec = n.children_by_field_name("type", &mut c).collect(); + for (i, ty) in types.into_iter().enumerate() { + if let Some(t) = chain_text(self.text(ty)) { + self.emit(RefKind::Field, ty, &i.to_string(), t); + } } true } @@ -682,7 +776,10 @@ impl<'a> Walker<'a> { .map(|p| { let t = self.text(p); let last = t.rsplit("::").next().unwrap_or(t).trim(); - if last.chars().next().is_some_and(|c| c.is_ascii_uppercase()) { + // `Arc::new(x)` is typed by `x` (see `rust_chain`). + if last.chars().next().is_some_and(|c| c.is_ascii_uppercase()) + && !matches!(last, "Arc" | "Rc" | "Box") + { last } else { "" @@ -698,6 +795,121 @@ impl<'a> Walker<'a> { } } + /// A declared type (`let x: T`, a parameter): a plain type's name, or a typed chain + /// for a generic one (`Arc`, `Option<&Item>`), which the resolver unwraps. + fn rust_declared_type(&self, t: Node) -> String { + let mut inner = t; + while inner.kind() == "reference_type" { + match inner.child_by_field_name("type") { + Some(x) => inner = x, + None => break, + } + } + match inner.kind() { + "generic_type" | "tuple_type" => chain_text(self.text(inner)) + .map(|c| format!("~T{c}")) + .unwrap_or_default(), + _ => rust_type_leaf(inner) + .map(|l| self.text(l).to_string()) + .unwrap_or_default(), + } + } + + /// The steps that type a Rust expression (see the module docs), or `None` when some + /// part of it is not something we follow (a literal, a closure, an index, a macro). + fn rust_chain(&self, v: Node, depth: u8) -> Option> { + if depth > MAX_CHAIN_STEPS as u8 { + return None; + } + let mut out = match v.kind() { + "self" => vec![format!("T{}", self.owner()?)], + "identifier" => { + let ty = self.local_type(self.text(v))?; + if let Some(chain) = ty.strip_prefix('~') { + chain.split('|').map(str::to_string).collect() + } else if ty.is_empty() { + return None; + } else { + vec![format!("T{ty}")] + } + } + "parenthesized_expression" | "reference_expression" => { + let inner = v.named_child(v.named_child_count().checked_sub(1)? as u32)?; + return self.rust_chain(inner, depth + 1); + } + "try_expression" => { + let mut c = self.rust_chain(v.named_child(0)?, depth + 1)?; + c.push("?".into()); + c + } + "await_expression" => { + let mut c = self.rust_chain(v.named_child(0)?, depth + 1)?; + c.push("^".into()); + c + } + "field_expression" => { + let field = v.child_by_field_name("field")?; + let mut c = self.rust_chain(v.child_by_field_name("value")?, depth + 1)?; + c.push(format!(":{}", self.text(field))); + c + } + "struct_expression" => { + let name = v.child_by_field_name("name").and_then(rust_type_leaf)?; + vec![format!("T{}", self.text(name))] + } + "call_expression" => { + let mut f = v.child_by_field_name("function")?; + if f.kind() == "generic_function" { + f = f.child_by_field_name("function")?; + } + match f.kind() { + "identifier" => { + let name = self.text(f); + if RUST_SKIP_CALLS.contains(&name) { + return None; + } + vec![format!("F#{name}")] + } + "scoped_identifier" => { + let name = self.text(f.child_by_field_name("name")?); + let path = f.child_by_field_name("path").map(|p| self.text(p)); + // `Arc::new(x)`: method calls see `x` through the pointer. + if name == "new" && path.is_some_and(|p| matches!(p, "Arc" | "Rc" | "Box")) + { + let args = v.child_by_field_name("arguments")?; + if args.named_child_count() != 1 { + return None; + } + return self.rust_chain(args.named_child(0)?, depth + 1); + } + let q = path.map(|p| self.qualify_path(p)).unwrap_or_default(); + vec![format!("F{q}#{name}")] + } + "field_expression" => { + let field = f.child_by_field_name("field")?; + if field.kind() != "field_identifier" { + return None; + } + let mut c = self.rust_chain(f.child_by_field_name("value")?, depth + 1)?; + c.push(format!(".{}", self.text(field))); + c + } + _ => return None, + } + } + _ => return None, + }; + if out.len() > MAX_CHAIN_STEPS { + return None; + } + let len: usize = out.iter().map(|s| s.len() + 1).sum(); + if len > MAX_CHAIN_BYTES || out.iter().any(|s| s.contains(['|', '~'])) { + return None; + } + out.shrink_to_fit(); + Some(out) + } + fn rust_call(&mut self, f: Node) { match f.kind() { "identifier" => { @@ -1783,6 +1995,33 @@ fn strip_generics(s: &str) -> String { out } +/// The identifier a pattern binds when it is just a name (`x`, `mut x`, `ref x`). +fn plain_binding(p: Node) -> Option { + match p.kind() { + "identifier" => Some(p), + "mut_pattern" | "ref_pattern" => { + let id = p.named_child(p.named_child_count().checked_sub(1)? as u32)?; + (id.kind() == "identifier").then_some(id) + } + _ => None, + } +} + +/// A field of a `struct` (not of an enum variant). +fn in_struct(field: Node) -> bool { + field + .parent() + .and_then(|l| l.parent()) + .is_some_and(|p| p.kind() == "struct_item") +} + +/// Type text as a typed chain carries it (whitespace collapsed), or `None` when it is too long +/// or holds a character the chain encoding reserves. +fn chain_text(t: &str) -> Option { + let t = t.split_whitespace().collect::>().join(" "); + (t.len() <= MAX_CHAIN_BYTES && !t.contains(['|', '~'])).then_some(t) +} + /// Whitespace-free receiver text, capped at a char boundary. fn compact(s: &str) -> String { let mut out: String = s.chars().filter(|c| !c.is_whitespace()).collect(); diff --git a/crates/atlas-codeindex/src/import_resolve.rs b/crates/atlas-codeindex/src/import_resolve.rs index 6cee02453..21fc13177 100644 --- a/crates/atlas-codeindex/src/import_resolve.rs +++ b/crates/atlas-codeindex/src/import_resolve.rs @@ -12,7 +12,7 @@ use std::collections::{BTreeSet, HashMap, HashSet}; use std::path::{Path, PathBuf}; use std::rc::Rc; -use crate::rust_crates::{join_rel, module_tree, parent_dir, CrateGraph, Libs}; +use crate::rust_crates::{join_rel, module_tree_with_aliases, parent_dir, CrateGraph, Libs}; use crate::universe::{family, Universe, GO, PY, RUST, TS}; #[derive(Debug, Clone, Default, PartialEq, Eq)] @@ -64,16 +64,23 @@ impl ProjectConfig { } /// `files.module` for every file, aligned with `u.files`. -pub(crate) fn compute_modules(u: &Universe, cfg: &ProjectConfig) -> Vec { +/// Each file's module key, plus the extra keys of Rust files mounted by more than one +/// `mod` declaration (a `tests/support/mod.rs` every test crate declares): install those +/// with [`install_module_aliases`] once the modules are set. +pub(crate) fn compute_modules( + u: &Universe, + cfg: &ProjectConfig, +) -> (Vec, Vec<(String, i64)>) { let rust_files: std::collections::BTreeMap = u .files .iter() .filter(|f| family(&f.lang) == RUST) .map(|f| (f.rel.clone(), f.id)) .collect(); - let tree = module_tree(&cfg.crates, &rust_files, &u.mods); + let (tree, aliases) = module_tree_with_aliases(&cfg.crates, &rust_files, &u.mods); let init_dirs = py_init_dirs(u); - u.files + let modules = u + .files .iter() .map(|f| match family(&f.lang) { RUST => tree @@ -85,7 +92,19 @@ pub(crate) fn compute_modules(u: &Universe, cfg: &ProjectConfig) -> Vec GO => go_module(&f.rel), _ => f.rel.clone(), }) - .collect() + .collect(); + (modules, aliases) +} + +/// Make each alias module key list its file too, so paths through it (`support::helper()`, +/// `use support::helper`) resolve in every crate that mounts the file. +pub(crate) fn install_module_aliases(u: &mut Universe, aliases: &[(String, i64)]) { + for (module, fid) in aliases { + let files = u.by_module.entry((RUST, module.clone())).or_default(); + if let Err(at) = files.binary_search(fid) { + files.insert(at, *fid); + } + } } fn ts_module(rel: &str) -> String { diff --git a/crates/atlas-codeindex/src/lib.rs b/crates/atlas-codeindex/src/lib.rs index 87493f509..d2f2ecb0a 100644 --- a/crates/atlas-codeindex/src/lib.rs +++ b/crates/atlas-codeindex/src/lib.rs @@ -32,6 +32,7 @@ mod query; mod repomap; mod resolve; mod rust_crates; +mod rust_types; mod scan; mod schema_v2; mod schema_v3; @@ -60,7 +61,7 @@ pub use lang::Lang; pub use query::{is_within, SymbolHit, SymbolQuery, SymbolSource}; pub use repomap::RepoMapFocus; pub use semantic::{ChunkHit, SemanticQuery}; -pub use skip::{SkipReason, MAX_FILE_BYTES}; +pub use skip::{skipped_by_dir, SkipReason, MAX_FILE_BYTES}; pub use store::{split_name, EXTRACTOR_VERSION, SCHEMA_VERSION}; pub use update::{BuildProgress, BuildStats, IndexStatus, UpdateStats}; pub use vectors::{Embedder, VectorStats}; diff --git a/crates/atlas-codeindex/src/query.rs b/crates/atlas-codeindex/src/query.rs index 2e5142553..a47fb6136 100644 --- a/crates/atlas-codeindex/src/query.rs +++ b/crates/atlas-codeindex/src/query.rs @@ -165,6 +165,17 @@ impl CodeIndex { self.find_symbol_where(q, &|_| true) } + /// [`CodeIndex::find_symbol`] plus how many of all the matches (every + /// page, after filters) are the name asked for: its exact name or + /// qualified name (or a qualified suffix), or the name ignoring case. + /// 0 means every hit is a prefix or word match, never the symbol itself. + pub fn find_symbol_with_exact( + &self, + q: &SymbolQuery, + ) -> Result<(Vec, Page, usize), IndexError> { + self.find_symbol_ranked(q, &|_| true) + } + /// [`CodeIndex::find_symbol`] keeping only the hits `keep` accepts, filtered /// with the query's own filters, before ranking and paging. pub(crate) fn find_symbol_where( @@ -172,6 +183,15 @@ impl CodeIndex { q: &SymbolQuery, keep: &dyn Fn(&SymbolHit) -> bool, ) -> Result<(Vec, Page), IndexError> { + self.find_symbol_ranked(q, keep) + .map(|(hits, page, _)| (hits, page)) + } + + fn find_symbol_ranked( + &self, + q: &SymbolQuery, + keep: &dyn Fn(&SymbolHit) -> bool, + ) -> Result<(Vec, Page, usize), IndexError> { let query = q.query.trim(); let Some(expr) = match_expr(query) else { return Err(IndexError::Invalid("query has no letters or digits".into())); @@ -205,6 +225,10 @@ impl CodeIndex { .then(a.id.cmp(&b.id)) }); let total = scored.len(); + let exact = scored + .iter() + .take_while(|(h, _)| tier(h, query) <= 1) + .count(); let page_hits: Vec = scored .into_iter() .skip(q.offset) @@ -220,7 +244,7 @@ impl CodeIndex { next_offset, truncation: next_offset.map(|_| "page_limit"), }; - Ok((page_hits, page)) + Ok((page_hits, page, exact)) } /// Every symbol in one file, in source order (parents before members). diff --git a/crates/atlas-codeindex/src/resolve.rs b/crates/atlas-codeindex/src/resolve.rs index 57c69ec25..7c720bb46 100644 --- a/crates/atlas-codeindex/src/resolve.rs +++ b/crates/atlas-codeindex/src/resolve.rs @@ -5,7 +5,14 @@ //! paths) → `import_map` 0.95 → `import_map_suffix` 0.85 (incl. glob imports) → `same_module` //! 0.90 → candidates by simple name (none if more than 256) → `qualified_suffix` 0.90 → //! `unique_name` 0.75 (×0.5 when not import-reachable) → `suffix_match` 0.55·min(1, 3/n) -//! (×0.5 when nothing is import-reachable). +//! (×0.5 when nothing is import-reachable). A Rust glob import yields to an item of the +//! module itself (`use demo::*; fn id()`), and a bare Rust name never binds a method. +//! Before all of it, a Rust receiver typed by a chain (`~…~…`, see `graph_extract`) binds +//! `typed_receiver` 0.90: the chain is evaluated through return types, struct fields, `?`, +//! `.await`, tuple elements and modelled std calls (`unwrap`, `lock`, auto-deref through +//! `Arc`/`Box`/guards) to a project type, whose impl must have the method (an impl whose +//! own type resolves elsewhere never counts; one that does not resolve, in another crate, +//! binds at 0.375); otherwise the receiver falls back to its text. //! Guards: same language family only; `receiver_chain_admits`; weak member calls (`x.f()` on an //! untyped receiver) never bind by name alone in Python/JS/TS, and in Rust/Go only to an //! import-reachable method (an unreachable one is mostly std's: `path.display()`); @@ -23,6 +30,7 @@ use std::collections::{HashMap, HashSet}; use crate::graph_extract::RefKind; use crate::import_resolve::{rust_abs, rust_longest_module, ImportTarget, ReexportGraph}; use crate::rust_crates::Libs; +use crate::rust_types::{self, RType, DEREF_WRAPPERS, FALLIBLE, LOCKS}; use crate::universe::{SymRow, Universe, GO, PY, RUST, TS}; pub(crate) const MAX_CANDIDATES: usize = 256; @@ -31,6 +39,15 @@ pub(crate) const MAX_CANDIDATES: usize = 256; /// reaches, is a guess, and a guess at hop 1 would read as a direct caller. pub(crate) const MIN_CONFIDENCE: f32 = 0.3; +/// A receiver typed by its chain, bound to that type's own method. +const TYPED_RECEIVER: f32 = 0.9; +/// A typed receiver's method found only by its impl's name in another crate (the impl's +/// type did not resolve): `unique_name`'s unreachable score, a name guess no import backs. +const CROSS_CRATE_NAME_MATCH: f32 = 0.375; +/// The cascade's import- and path-backed strategies score at least this; below it a +/// target is a name guess. +const FIRM: f32 = 0.85; + pub(crate) const CALLABLE_KINDS: &[&str] = &["fn", "function", "method", "constructor", "macro", "func"]; /// Calls may also construct: `new Foo()`, `Foo()`, Rust tuple structs, Go conversions. @@ -175,6 +192,8 @@ pub(crate) struct Resolver<'u> { memo: HashMap<(i64, u8, String, String), Option>, /// A symbol no candidate may be (see [`Resolver::resolve`]). exclude: Option, + /// Impl/trait symbol id → the type definition it is for ([`Resolver::parent_type`]). + parent_ty: HashMap>, } impl<'u> Resolver<'u> { @@ -211,11 +230,15 @@ impl<'u> Resolver<'u> { reach: HashMap::new(), memo: HashMap::new(), exclude: None, + parent_ty: HashMap::new(), } } /// Resolve one reference to `(src symbol, target)`. `None` = no edge. pub(crate) fn resolve(&mut self, r: &RefRow) -> Option<(i64, Resolved)> { + if r.kind == RefKind::Field { + return None; + } let (src, receiver) = match r.kind { RefKind::Inherit | RefKind::Impl => { let (src_ty, qual) = r @@ -288,6 +311,15 @@ impl<'u> Resolver<'u> { } fn cascade(&mut self, f: i64, kind: RefKind, receiver: &str, name: &str) -> Option { + if let Some(rest) = receiver.strip_prefix('~') { + let (chain, fallback) = rest.split_once('~').unwrap_or((rest, "")); + if kind == RefKind::Call { + if let Some(hit) = self.typed_member(f, chain, name) { + return Some(hit); + } + } + return self.cascade(f, kind, fallback, name); + } let u = self.u; let fam = u.fam_of(f); let (segs, qualified) = parse_receiver(receiver); @@ -299,8 +331,15 @@ impl<'u> Resolver<'u> { return None; } let exclude = self.exclude; - let ok = - |s: &SymRow| kind_ok(kind, s) && u.fam_of(s.file_id) == fam && Some(s.id) != exclude; + // A bare Rust name (`id(1)`) is never an associated fn or method: those take a path + // (`Self::id`) or a receiver (`x.id()`). + let bare_rust = fam == RUST && segs.is_empty(); + let ok = |s: &SymRow| { + kind_ok(kind, s) + && u.fam_of(s.file_id) == fam + && Some(s.id) != exclude + && !(bare_rust && s.kind == "method") + }; let value = kind == RefKind::Value; // Rust paths through the real module tree. @@ -321,6 +360,7 @@ impl<'u> Resolver<'u> { } // S1 import_map / S1b import_map_suffix. + let mut glob_hit: Option = None; if let Some((b, rest)) = self.binding_for(f, &segs, name) { if b.external { return None; @@ -357,17 +397,10 @@ impl<'u> Resolver<'u> { } } } else if segs.is_empty() { - let globs: Vec = self - .bindings - .get(&f) - .map(|v| v.iter().filter(|b| b.glob).cloned().collect()) - .unwrap_or_default(); - for b in globs { - let (files, need) = self.binding_scope(fam, &b, &[]); - let need_refs: Vec<&str> = need.iter().map(String::as_str).collect(); - let c = self.in_files(&files, name, &ok, &need_refs, true); - if !c.is_empty() { - return Some(self.pick(f, &c, 0.85, "import_map_suffix")); + if let Some(hit) = self.glob_import(f, fam, name, &ok) { + glob_hit = Some(hit); + if fam != RUST { + return glob_hit; } } } @@ -381,9 +414,19 @@ impl<'u> Resolver<'u> { .filter(|i| admits(u, &segs, *i) && (!member || is_method(u, *i))) .collect(); if !c.is_empty() { - return Some(self.pick(f, &c, 0.90, "same_module")); + let local = self.pick(f, &c, 0.90, "same_module"); + // A Rust item of the module shadows a glob import (`use demo::*; fn id()`); + // when both name the same item (`use super::*` in an inline `mod tests`), + // the glob is how it got there. + return Some(match glob_hit { + Some(g) if g.dst == local.dst => g, + _ => local, + }); } } + if glob_hit.is_some() { + return glob_hit; + } // S3: by simple name. let all: Vec = self @@ -447,6 +490,307 @@ impl<'u> Resolver<'u> { Some(self.pick(f, &pool, conf, "suffix_match")) } + /// S1b through a glob import (`use a::*`, `from a import *`). + fn glob_import( + &mut self, + f: i64, + fam: u8, + name: &str, + ok: &dyn Fn(&SymRow) -> bool, + ) -> Option { + let globs: Vec = self + .bindings + .get(&f) + .map(|v| v.iter().filter(|b| b.glob).cloned().collect()) + .unwrap_or_default(); + for b in globs { + let (files, need) = self.binding_scope(fam, &b, &[]); + let need_refs: Vec<&str> = need.iter().map(String::as_str).collect(); + let c = self.in_files(&files, name, ok, &need_refs, true); + if !c.is_empty() { + return Some(self.pick(f, &c, 0.85, "import_map_suffix")); + } + } + None + } + + /// `x.name()` on a receiver typed by a chain (see `graph_extract`): the method of the + /// type the chain ends at. `None` = the chain does not type, or the type has no such + /// project method (a std or trait method); the caller then reads the receiver as text. + fn typed_member(&mut self, f: i64, chain: &str, name: &str) -> Option { + let (t, home) = self.chain_type(f, chain)?; + match self.member(t, home, name)? { + Member::Method(i, confidence) => Some(Resolved { + dst: self.u.syms[i].id, + confidence, + strategy: "typed_receiver", + }), + Member::Std(_) => None, + } + } + + /// The type a chain evaluates to, with the file its text was written in. + fn chain_type(&mut self, f: i64, chain: &str) -> Option<(RType, i64)> { + let mut cur: Option<(RType, i64)> = None; + for step in chain.split('|') { + let mut chars = step.chars(); + let op = chars.next()?; + let arg = chars.as_str(); + cur = Some(match op { + 'T' => (rust_types::parse(arg)?, f), + 'F' => { + let (q, n) = arg.rsplit_once('#')?; + let d = self.cascade(f, RefKind::Call, q, n)?; + self.result_of(*self.u.sym_ix.get(&d.dst)?)? + } + '.' => { + let (t, home) = cur?; + match self.member(t, home, arg)? { + Member::Std(t) => (t, home), + Member::Method(i, _) => self.result_of(i)?, + } + } + ':' => { + let (t, home) = cur?; + self.field_type(t, home, arg)? + } + '?' => { + let (t, home) = cur?; + let t = deref(t); + if !FALLIBLE.contains(&t.name()) { + return None; + } + (t.arg(0)?, home) + } + '^' => { + let (t, home) = cur?; + // An `async fn`'s declared type is already what `.await` yields. + if t.name() == "Future" || t.name() == "IntoFuture" { + (t.arg(0)?, home) + } else { + (t, home) + } + } + '#' => { + let (t, home) = cur?; + let t = deref(t); + if !t.is_tuple() { + return None; + } + (t.arg(arg.parse().ok()?)?, home) + } + _ => return None, + }); + } + cur + } + + /// What `.name()` on `t` is: a std operation we model (`unwrap`, `lock`, `clone`, + /// auto-deref through `Arc`/`Box`/guards), or the type's own method. + fn member(&mut self, mut t: RType, home: i64, name: &str) -> Option { + for _ in 0..8 { + let n = t.name(); + let unwrap = matches!( + name, + "unwrap" | "expect" | "unwrap_or_default" | "unwrap_or" | "unwrap_or_else" + ); + if unwrap && FALLIBLE.contains(&n) { + return Some(Member::Std(t.arg(0)?)); + } + if matches!( + name, + "lock" + | "read" + | "write" + | "try_lock" + | "blocking_lock" + | "blocking_read" + | "blocking_write" + ) && LOCKS.contains(&n) + { + return Some(Member::Std(t.arg(0)?)); + } + if DEREF_WRAPPERS.contains(&n) { + t = t.arg(0)?; + continue; + } + if matches!( + name, + "clone" + | "to_owned" + | "as_ref" + | "as_mut" + | "borrow" + | "borrow_mut" + | "deref" + | "deref_mut" + ) || unwrap + { + // `guard.unwrap()` on a lock's result we already unwrapped, `x.clone()`. + return Some(Member::Std(t)); + } + let (t, home) = self.expand_alias(t, home)?; + if FALLIBLE.contains(&t.name()) || t.is_tuple() { + return None; + } + return self + .method_of(&t, home, name) + .map(|(i, c)| Member::Method(i, c)); + } + None + } + + /// A project method `name` of type `t` (`t` written in `home`): one whose parent impl or + /// trait is for that very type. A parent whose type resolves (from the impl's own file) + /// to another definition is never one, however it is named: `Foo` in a crate that has + /// a `Foo` of its own is that crate's. A parent whose type does not resolve falls back + /// to its name: in the type's own crate at the typed confidence, in another crate at + /// [`CROSS_CRATE_NAME_MATCH`]. + fn method_of(&mut self, t: &RType, home: i64, name: &str) -> Option<(usize, f32)> { + let u = self.u; + // A type no project symbol answers to (`tempfile::TempDir`, `std::path::Path`) has + // none of the project's methods, whatever their parent is named. + let ti = self.type_sym(t, home)?; + let tid = u.syms[ti].id; + let tname = u.syms[ti].name.as_str(); + let named: Vec<(usize, usize)> = self + .by_name(name) + .iter() + .copied() + .filter_map(|i| { + let s = &u.syms[i]; + let pi = *s.parent.and_then(|p| u.sym_ix.get(&p))?; + (CALLABLE_KINDS.contains(&s.kind.as_str()) + && u.fam_of(s.file_id) == RUST + && Some(s.id) != self.exclude + && u.syms[pi].name == tname) + .then_some((i, pi)) + }) + .collect(); + let mut exact = Vec::new(); + let mut by_name = Vec::new(); + for (i, pi) in named { + match self.parent_type(pi) { + Some(id) if id == tid => exact.push(i), + Some(_) => {} + None => by_name.push(i), + } + } + let home_crate = crate_key(&u.file(u.syms[ti].file_id).module); + let in_home = |i: &usize| crate_key(&u.file(u.syms[*i].file_id).module) == home_crate; + // The type's own crate first: its inherent methods shadow another crate's trait + // impl of the same name. + let (exact_home, exact_other): (Vec, Vec) = + exact.into_iter().partition(in_home); + let (own, other): (Vec, Vec) = by_name.into_iter().partition(in_home); + let (cands, confidence) = if !exact_home.is_empty() { + (exact_home, TYPED_RECEIVER) + } else if !exact_other.is_empty() { + (exact_other, TYPED_RECEIVER) + } else if !own.is_empty() { + (own, TYPED_RECEIVER) + } else if !other.is_empty() { + (other, CROSS_CRATE_NAME_MATCH) + } else { + return None; + }; + let f = u.syms[cands[0]].file_id; + let d = self.pick(f, &cands, confidence, "typed_receiver"); + Some((*u.sym_ix.get(&d.dst)?, confidence)) + } + + /// The type definition an impl, trait or type symbol `pi` is for, resolved from its own + /// file; `None` when nothing resolves it firmly (a name-only guess is no answer). + fn parent_type(&mut self, pi: usize) -> Option { + let p = &self.u.syms[pi]; + if TYPE_KINDS.contains(&p.kind.as_str()) { + return Some(p.id); + } + if let Some(hit) = self.parent_ty.get(&p.id) { + return *hit; + } + let (id, file, name) = (p.id, p.file_id, p.name.clone()); + let exclude = self.exclude.take(); + let res = self + .cascade(file, RefKind::Type, "", &name) + .filter(|d| d.confidence >= FIRM) + .map(|d| d.dst); + self.exclude = exclude; + self.parent_ty.insert(id, res); + res + } + + /// The type symbol `t` names, as seen from `home`. + fn type_sym(&mut self, t: &RType, home: i64) -> Option { + if t.is_tuple() || t.name().is_empty() { + return None; + } + let d = self.cascade(home, RefKind::Type, &t.qualifier(), t.name())?; + self.u.sym_ix.get(&d.dst).copied() + } + + /// `type Store = Arc;` → `Arc`, as many times as it takes (bounded). + fn expand_alias(&mut self, mut t: RType, mut home: i64) -> Option<(RType, i64)> { + for _ in 0..4 { + let Some(i) = self.type_sym(&t, home) else { + return Some((t, home)); + }; + let s = &self.u.syms[i]; + if s.kind != "type" && s.kind != "type_alias" { + return Some((t, home)); + } + let (_, rhs) = s.sig.split_once('=')?; + let next = rust_types::parse(rhs.trim().trim_end_matches(';'))?; + if DEREF_WRAPPERS.contains(&next.name()) || LOCKS.contains(&next.name()) { + return Some((next, s.file_id)); + } + home = s.file_id; + t = next; + } + Some((t, home)) + } + + /// The value a call to symbol `i` produces: a struct's constructor makes the struct, a + /// fn its declared return type with `Self` read as the impl type. + fn result_of(&mut self, i: usize) -> Option<(RType, i64)> { + let s = &self.u.syms[i]; + if CONSTRUCTIBLE_KINDS.contains(&s.kind.as_str()) { + return Some(( + RType { + path: vec![s.name.clone()], + args: Vec::new(), + }, + s.file_id, + )); + } + let ret = rust_types::parse(rust_types::return_type(&s.sig)?)?; + let owner = self.owner_name(i); + Some((with_self(ret, owner.as_deref()), s.file_id)) + } + + /// The type a field of `t` holds. + fn field_type(&mut self, t: RType, home: i64, field: &str) -> Option<(RType, i64)> { + let t = deref(t); + if t.is_tuple() { + return Some((t.arg(field.parse().ok()?)?, home)); + } + let (t, home) = self.expand_alias(t, home)?; + let t = deref(t); + let si = self.type_sym(&t, home)?; + let s = &self.u.syms[si]; + let (_, text) = self.u.fields.get(&s.id)?.iter().find(|(n, _)| n == field)?; + let ft = rust_types::parse(text)?; + Some((with_self(ft, Some(&s.name)), s.file_id)) + } + + /// The impl or trait a method belongs to, by name. + fn owner_name(&self, i: usize) -> Option { + let s = &self.u.syms[i]; + let pi = s.parent.and_then(|p| self.u.sym_ix.get(&p))?; + let p = &self.u.syms[*pi]; + matches!(p.kind.as_str(), "impl" | "trait" | "struct" | "enum").then(|| p.name.clone()) + } + fn by_name(&self, name: &str) -> &'u [usize] { self.u.by_name.get(name).map(Vec::as_slice).unwrap_or(&[]) } @@ -633,6 +977,43 @@ impl<'u> Resolver<'u> { } } +enum Member { + /// A std operation's result type. + Std(RType), + /// A project method (symbol index), with how sure the binding is. + Method(usize, f32), +} + +/// Through `Arc`, `Box`, guards and the like to the type a method call sees. +fn deref(mut t: RType) -> RType { + for _ in 0..8 { + if !DEREF_WRAPPERS.contains(&t.name()) { + break; + } + match t.arg(0) { + Some(inner) => t = inner, + None => break, + } + } + t +} + +/// `Self` → the impl type, anywhere in `t`. +fn with_self(mut t: RType, owner: Option<&str>) -> RType { + if let Some(o) = owner { + if t.path.len() == 1 && t.path[0] == "Self" { + t.path[0] = o.to_string(); + } + t.args = t.args.into_iter().map(|a| with_self(a, owner)).collect(); + } + t +} + +/// The crate part of a module key (`atlas_memory::record` → `atlas_memory`). +fn crate_key(module: &str) -> &str { + module.split("::").next().unwrap_or(module) +} + /// `""` → no receiver; `"a::B::"` → qualified segments; `"x.y"` → member segments. pub(crate) fn parse_receiver(receiver: &str) -> (Vec<&str>, bool) { if receiver.is_empty() { @@ -653,6 +1034,7 @@ fn kind_ok(kind: RefKind, s: &SymRow) -> bool { RefKind::Call => CALLABLE_KINDS.contains(&k) || CONSTRUCTIBLE_KINDS.contains(&k), RefKind::Value => CALLABLE_KINDS.contains(&k), RefKind::Type | RefKind::Inherit | RefKind::Impl => TYPE_KINDS.contains(&k), + RefKind::Field => false, } } @@ -787,6 +1169,7 @@ mod tests { qn: qn.into(), exported: true, is_test: false, + sig: String::new(), } } @@ -1180,6 +1563,61 @@ mod tests { ); } + /// Two crates each define a `Foo`. A receiver typed as alpha's `Foo` binds a method only + /// through an impl for alpha's `Foo`: never beta's `impl Foo`, which is for beta's own. + #[test] + fn a_typed_receiver_binds_its_own_types_impl_not_a_same_named_one_elsewhere() { + let u = Universe::from_parts( + vec![ + file(1, "alpha/src/lib.rs", "rust", "alpha"), + file(2, "beta/src/lib.rs", "rust", "beta"), + file(3, "beta/src/show.rs", "rust", "beta::show"), + file(4, "gamma/src/lib.rs", "rust", "gamma"), + ], + vec![ + sym(1, 1, None, "fn", "caller", "caller"), + sym(10, 1, None, "struct", "Foo", "Foo"), + sym(11, 1, None, "impl", "Foo", "Foo"), + sym(12, 1, Some(11), "method", "get", "Foo::get"), + // beta's own `Foo`, and its inherent `run`. + sym(20, 2, None, "struct", "Foo", "Foo"), + sym(21, 2, None, "impl", "Foo", "Foo"), + sym(22, 2, Some(21), "method", "run", "Foo::run"), + // `use alpha::Foo; impl Show for Foo`: a trait impl for alpha's `Foo`. + sym(30, 3, None, "impl", "Foo", "Foo"), + sym(31, 3, Some(30), "method", "show", "Foo::show"), + sym(32, 3, Some(30), "method", "get", "Foo::get"), + // An impl whose `Foo` resolves to nothing firm: a name guess. + sym(40, 4, None, "impl", "Foo", "Foo"), + sym(41, 4, Some(40), "method", "size", "Foo::size"), + ], + vec![imp(3, "Foo", "alpha", "Foo")], + HashMap::new(), + ); + let t = vec![ImportTarget { + file: Some(1), + module: Some("alpha".into()), + symbol: Some("Foo".into()), + glob: false, + }]; + let mut r = resolver(&u, t); + let got = r.resolve(&call(1, 1, "run", "~TFoo~x")); + assert!(got.is_none_or(|(_, d)| d.dst != 22), "{got:?}"); + let (_, d) = r.resolve(&call(1, 1, "show", "~TFoo~x")).unwrap(); + assert_eq!( + (d.dst, d.strategy, d.confidence), + (31, "typed_receiver", 0.9) + ); + // alpha's inherent `get` shadows beta's trait impl of the same name. + let (_, d) = r.resolve(&call(3, 31, "get", "~TFoo~x")).unwrap(); + assert_eq!(d.dst, 12); + let (_, d) = r.resolve(&call(1, 1, "size", "~TFoo~x")).unwrap(); + assert_eq!( + (d.dst, d.strategy, d.confidence), + (41, "typed_receiver", CROSS_CRATE_NAME_MATCH) + ); + } + #[test] fn impl_refs_resolve_both_ends() { let u = Universe::from_parts( diff --git a/crates/atlas-codeindex/src/rust_crates.rs b/crates/atlas-codeindex/src/rust_crates.rs index 704935dfa..766adf311 100644 --- a/crates/atlas-codeindex/src/rust_crates.rs +++ b/crates/atlas-codeindex/src/rust_crates.rs @@ -467,13 +467,32 @@ pub(crate) fn join_rel(base: &str, tail: &str) -> Option { /// Module path for every Rust file reachable from a crate root through `mod` declarations. /// `files`: rel → file id (Rust files only). `mods`: file id → its `mod x;` declarations. +#[cfg(test)] pub fn module_tree( graph: &CrateGraph, files: &BTreeMap, mods: &HashMap>, ) -> BTreeMap { + module_tree_with_aliases(graph, files, mods).0 +} + +/// Most alias module keys one project may add (a guard against `mod` cycles). +const MAX_MODULE_ALIASES: usize = 20_000; + +/// Each file's module path through the `mod` tree (the first crate to reach it), plus every +/// other module path a file is also mounted at: a `tests/support/mod.rs` that each test +/// crate declares with `mod support;` is `handshake#bin::support` in the map and also +/// `invariants#bin::support`, `manager#bin::support`, … in the aliases, with its own `mod` +/// children under each. +pub fn module_tree_with_aliases( + graph: &CrateGraph, + files: &BTreeMap, + mods: &HashMap>, +) -> (BTreeMap, Vec<(String, i64)>) { let rel_of: HashMap = files.iter().map(|(r, id)| (*id, r.as_str())).collect(); let mut assigned: BTreeMap = BTreeMap::new(); + let mut aliases: Vec<(String, i64)> = Vec::new(); + let mut seen_alias: BTreeSet<(i64, String)> = BTreeSet::new(); let mut queue: VecDeque<(i64, String)> = VecDeque::new(); for root in &graph.roots { if let Some(&id) = files.get(&root.root_rel) { @@ -527,19 +546,29 @@ pub fn module_tree( let Some(child) = candidates.iter().find_map(|c| files.get(c)) else { continue; }; - if assigned.contains_key(child) { - continue; - } let mut child_mod = module.clone(); for seg in inline.iter().chain(std::iter::once(&m.name.as_str())) { child_mod.push_str("::"); child_mod.push_str(seg); } - assigned.insert(*child, child_mod.clone()); - queue.push_back((*child, child_mod)); + match assigned.get(child) { + None => { + assigned.insert(*child, child_mod.clone()); + queue.push_back((*child, child_mod)); + } + Some(primary) if *primary != child_mod => { + if aliases.len() < MAX_MODULE_ALIASES + && seen_alias.insert((*child, child_mod.clone())) + { + aliases.push((child_mod.clone(), *child)); + queue.push_back((*child, child_mod)); + } + } + Some(_) => {} + } } } - assigned + (assigned, aliases) } #[cfg(test)] diff --git a/crates/atlas-codeindex/src/rust_types.rs b/crates/atlas-codeindex/src/rust_types.rs new file mode 100644 index 000000000..a270fd93f --- /dev/null +++ b/crates/atlas-codeindex/src/rust_types.rs @@ -0,0 +1,305 @@ +//! Rust type text, read well enough to type a method receiver: `Arc>`, +//! `anyhow::Result<(Store, u32)>`, `Option<&Item>`, `impl Future`. Paths keep +//! their segments, generic arguments nest, references, lifetimes, `mut`, `dyn` and extra +//! `+` bounds are dropped. Anything else (fn pointers, arrays' lengths, macros) parses as far +//! as it can and is not followed further. + +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub(crate) struct RType { + /// Path segments, `[]` for a tuple. + pub path: Vec, + /// Generic arguments, or a tuple's elements. + pub args: Vec, +} + +impl RType { + pub(crate) fn name(&self) -> &str { + self.path.last().map(String::as_str).unwrap_or("") + } + + pub(crate) fn is_tuple(&self) -> bool { + self.path.is_empty() + } + + /// `a::b::` for `a::b::C`, `""` for `C`. + pub(crate) fn qualifier(&self) -> String { + let n = self.path.len().saturating_sub(1); + self.path[..n].iter().map(|s| format!("{s}::")).collect() + } + + /// The first generic argument / tuple element. + pub(crate) fn arg(&self, i: usize) -> Option { + self.args.get(i).cloned() + } +} + +/// Types a method call sees through (auto-deref): `Arc` has `T`'s methods. +pub(crate) const DEREF_WRAPPERS: &[&str] = &[ + "Arc", + "Rc", + "Box", + "Cow", + "Pin", + "MutexGuard", + "RwLockReadGuard", + "RwLockWriteGuard", + "Ref", + "RefMut", + "OwnedMutexGuard", + "ArcSwap", +]; + +/// `.lock()` / `.read()` / `.write()` hand out their argument (through a guard). +pub(crate) const LOCKS: &[&str] = &["Mutex", "RwLock", "ReentrantMutex"]; + +/// Types `?`, `.unwrap()` and `.expect()` take apart. +pub(crate) const FALLIBLE: &[&str] = &["Option", "Result", "LockResult", "Poll"]; + +pub(crate) fn parse(text: &str) -> Option { + let mut p = Parser { + s: text.as_bytes(), + i: 0, + depth: 0, + }; + p.ty() +} + +/// The declared return type of a fn signature (`… -> T`, before any `where`), or `None`. +pub(crate) fn return_type(sig: &str) -> Option<&str> { + let b = sig.as_bytes(); + let (mut paren, mut angle) = (0i32, 0i32); + let mut i = 0; + while i + 1 < b.len() { + match b[i] { + b'(' | b'[' => paren += 1, + b')' | b']' => paren -= 1, + b'<' => angle += 1, + b'>' if i > 0 && b[i - 1] != b'-' => angle -= 1, + b'-' if b[i + 1] == b'>' && paren == 0 && angle == 0 => { + let rest = sig[i + 2..].trim(); + let rest = rest.split(" where ").next().unwrap_or(rest); + let rest = rest.trim_end_matches('…').trim(); + return (!rest.is_empty()).then_some(rest); + } + _ => {} + } + i += 1; + } + None +} + +struct Parser<'a> { + s: &'a [u8], + i: usize, + depth: u8, +} + +impl Parser<'_> { + fn ws(&mut self) { + while self.i < self.s.len() && self.s[self.i].is_ascii_whitespace() { + self.i += 1; + } + } + + fn peek(&mut self) -> Option { + self.ws(); + self.s.get(self.i).copied() + } + + fn eat(&mut self, c: u8) -> bool { + if self.peek() == Some(c) { + self.i += 1; + true + } else { + false + } + } + + fn ident(&mut self) -> Option<&str> { + self.ws(); + let start = self.i; + while self.i < self.s.len() + && (self.s[self.i].is_ascii_alphanumeric() || self.s[self.i] == b'_') + { + self.i += 1; + } + (self.i > start).then(|| std::str::from_utf8(&self.s[start..self.i]).unwrap_or("")) + } + + fn keyword(&mut self, kw: &str) -> bool { + self.ws(); + let rest = &self.s[self.i..]; + if rest.starts_with(kw.as_bytes()) + && rest + .get(kw.len()) + .is_none_or(|c| !(c.is_ascii_alphanumeric() || *c == b'_')) + { + self.i += kw.len(); + true + } else { + false + } + } + + fn ty(&mut self) -> Option { + self.depth += 1; + if self.depth > 24 { + return None; + } + let out = self.ty_inner(); + self.depth -= 1; + out + } + + fn ty_inner(&mut self) -> Option { + while self.eat(b'&') { + if self.peek() == Some(b'\'') { + self.i += 1; + self.ident(); + } + self.keyword("mut"); + } + if self.keyword("dyn") || self.keyword("impl") { + let t = self.ty()?; + while self.eat(b'+') { + if self.peek() == Some(b'\'') { + self.i += 1; + self.ident(); + } else { + self.ty()?; + } + } + return Some(t); + } + if self.eat(b'(') { + let mut args = Vec::new(); + while !self.eat(b')') { + args.push(self.ty()?); + if !self.eat(b',') && self.peek() != Some(b')') { + return None; + } + } + return Some(RType { + path: Vec::new(), + args, + }); + } + if self.eat(b'[') { + let inner = self.ty()?; + while self.peek().is_some_and(|c| c != b']') { + self.i += 1; + } + self.eat(b']'); + return Some(RType { + path: vec!["[]".into()], + args: vec![inner], + }); + } + let mut path = Vec::new(); + let mut args = Vec::new(); + loop { + path.push(self.ident()?.to_string()); + self.ws(); + if self.s[self.i..].starts_with(b"::") { + self.i += 2; + if self.eat(b'<') { + args = self.generic_args()?; + break; + } + continue; + } + if self.eat(b'<') { + args = self.generic_args()?; + } + break; + } + Some(RType { path, args }) + } + + /// After `<`: types, lifetimes skipped, `Name = T` bindings kept as `T`. + fn generic_args(&mut self) -> Option> { + let mut args = Vec::new(); + while !self.eat(b'>') { + if self.peek() == Some(b'\'') { + self.i += 1; + self.ident(); + } else { + let save = self.i; + let binding = self.ident().is_some() && self.peek() == Some(b'='); + if binding { + self.i += 1; + } else { + self.i = save; + } + args.push(self.ty()?); + } + if !self.eat(b',') && self.peek() != Some(b'>') { + return None; + } + } + Some(args) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn show(t: &RType) -> String { + let head = if t.is_tuple() { + String::new() + } else { + t.path.join("::") + }; + if t.args.is_empty() && !t.is_tuple() { + return head; + } + let inner: Vec = t.args.iter().map(show).collect(); + if t.is_tuple() { + format!("({})", inner.join(",")) + } else { + format!("{head}<{}>", inner.join(",")) + } + } + + #[test] + fn types_parse_to_paths_and_arguments() { + let p = |s: &str| parse(s).map(|t| show(&t)); + assert_eq!(p("Arc>").as_deref(), Some("Arc>")); + assert_eq!( + p("&'a mut crate::store::Store").as_deref(), + Some("crate::store::Store") + ); + assert_eq!( + p("anyhow::Result<(Store, u32)>").as_deref(), + Some("anyhow::Result<(Store,u32)>") + ); + assert_eq!(p("Option<&Item>").as_deref(), Some("Option")); + assert_eq!( + p("impl std::future::Future + Send + 'static").as_deref(), + Some("std::future::Future") + ); + assert_eq!( + p("HashMap>").as_deref(), + Some("HashMap>") + ); + assert_eq!(p("Vec::").as_deref(), Some("Vec")); + } + + #[test] + fn a_return_type_is_after_the_top_level_arrow() { + assert_eq!( + return_type("pub fn open() -> std::io::Result"), + Some("std::io::Result") + ); + assert_eq!( + return_type("fn f(cb: impl Fn() -> u8) -> Out where T: X"), + Some("Out") + ); + assert_eq!(return_type("fn g(&self)"), None); + assert_eq!( + return_type("pub fn r(&self) -> Result< ( Id, u8, ), E, >"), + Some("Result< ( Id, u8, ), E, >") + ); + } +} diff --git a/crates/atlas-codeindex/src/scan.rs b/crates/atlas-codeindex/src/scan.rs index ff180bd3a..9cb7d04ee 100644 --- a/crates/atlas-codeindex/src/scan.rs +++ b/crates/atlas-codeindex/src/scan.rs @@ -7,7 +7,7 @@ //! back in path order before anything is written. use std::path::{Path, PathBuf}; -use std::sync::{mpsc, Arc, OnceLock}; +use std::sync::{mpsc, Arc, Mutex, OnceLock}; use std::time::{Duration, Instant, UNIX_EPOCH}; use atlas_search::CancelToken; @@ -63,10 +63,16 @@ pub(crate) fn mtime_ns(md: &std::fs::Metadata) -> i64 { .map_or(0, |d| i64::try_from(d.as_nanos()).unwrap_or(i64::MAX)) } +/// Directories a walk pruned, with the skip reason their files report. +type Pruned = Arc>>; + /// Gitignore-respecting walk settings shared by every walk the index does. pub(crate) fn walker(dir: &Path, root: &Path, rules: Arc) -> WalkBuilder { + walker_noting(dir, root, rules, None) +} + +fn base_walker(dir: &Path) -> WalkBuilder { let mut b = WalkBuilder::new(dir); - let root = root.to_path_buf(); b.hidden(false) .git_ignore(true) .git_global(true) @@ -74,17 +80,103 @@ pub(crate) fn walker(dir: &Path, root: &Path, rules: Arc) -> WalkBuilder .ignore(true) .parents(true) .require_git(false) - .follow_links(false) - .filter_entry(move |e| { - if !e.file_type().is_some_and(|t| t.is_dir()) || e.depth() == 0 { - return true; + .follow_links(false); + b +} + +/// [`walker`], noting each directory it prunes in `pruned` (when given). +fn walker_noting( + dir: &Path, + root: &Path, + rules: Arc, + pruned: Option, +) -> WalkBuilder { + let mut b = base_walker(dir); + let root = root.to_path_buf(); + b.filter_entry(move |e| { + if !e.file_type().is_some_and(|t| t.is_dir()) || e.depth() == 0 { + return true; + } + let Ok(rel) = e.path().strip_prefix(&root) else { + return true; + }; + let rel = rel_string(rel); + if !rules.prune_dir(&rel) { + return true; + } + if let (Some(pruned), Some(reason)) = (&pruned, rules.prune_reason(&rel)) { + pruned + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push((e.path().to_path_buf(), reason)); + } + false + }); + b +} + +/// Every source file under the pruned directories, with its directory's +/// reason. Pruning saves reading and parsing them, not knowing they exist: +/// without these the index would look complete while missing a whole tree. +/// Only names are looked at (no stat, no read); gitignore still applies. +fn pruned_files( + root: &Path, + pruned: Vec<(PathBuf, SkipReason)>, + cancel: &CancelToken, +) -> Vec<(String, SkipReason)> { + let Some(((first, _), rest)) = pruned.split_first() else { + return Vec::new(); + }; + let mut b = base_walker(first); + for (dir, _) in rest { + b.add(dir); + } + b.filter_entry(|e| { + e.depth() == 0 + || !e.file_type().is_some_and(|t| t.is_dir()) + || e.file_name().to_str().is_none_or(skip::counted_dir) + }); + let (tx, rx) = mpsc::channel::<(String, SkipReason)>(); + let pruned = Arc::new(pruned); + b.build_parallel().run(|| { + let tx = tx.clone(); + let pruned = pruned.clone(); + let cancel = cancel.clone(); + let root = root.to_path_buf(); + Box::new(move |entry| { + if cancel.is_cancelled() { + return WalkState::Quit; } - match e.path().strip_prefix(&root) { - Ok(rel) => !rules.prune_dir(&rel_string(rel)), - Err(_) => true, + let Ok(entry) = entry else { + return WalkState::Continue; + }; + if !entry.file_type().is_some_and(|t| t.is_file()) { + return WalkState::Continue; } - }); - b + let path = entry.path(); + let Ok(rel) = path.strip_prefix(&root).map(rel_string) else { + return WalkState::Continue; + }; + if Lang::from_path(&rel).is_none() { + return WalkState::Continue; + } + // The deepest pruned directory holding it (they never nest, but + // a walk root given twice must not double-count). + if let Some((_, reason)) = pruned + .iter() + .filter(|(d, _)| path.starts_with(d)) + .max_by_key(|(d, _)| d.as_os_str().len()) + { + let _ = tx.send((rel, *reason)); + } + WalkState::Continue + }) + }); + drop(tx); + let mut out: Vec<_> = rx.into_iter().collect(); + out.sort(); + out.dedup(); + out } /// Classify one walked or watched file by path and stat. @@ -111,8 +203,6 @@ pub(crate) fn classify( })) } -/// Every indexable file under `dir` (the root or a subdirectory), sorted by -/// path, plus the files skipped by path or size with their reasons. /// What a walk saw that is not a source file to index: one it skipped, or a /// graph config file (`Cargo.toml`, `tsconfig.json`, …) as a /// `rel size mtime_ns` stamp. @@ -121,8 +211,9 @@ enum Seen { Config(String), } -/// The indexable files under `dir`, the files skipped and why, and a stamp of -/// every graph config file met on the way (see [`config_stamp`]). +/// The indexable files under `dir`, sorted by path; every source file +/// skipped and why (by path, size, or under a pruned directory); and a stamp +/// of every graph config file met on the way (see [`config_stamp`]). pub(crate) fn discover( root: &Path, dir: &Path, @@ -130,41 +221,49 @@ pub(crate) fn discover( cancel: &CancelToken, ) -> (Vec, Vec<(String, SkipReason)>, Vec) { let (tx, rx) = mpsc::channel::(); - walker(dir, root, rules.clone()).build_parallel().run(|| { - let tx = tx.clone(); - let rules = rules.clone(); - let cancel = cancel.clone(); - Box::new(move |entry| { - if cancel.is_cancelled() { - return WalkState::Quit; - } - let Ok(entry) = entry else { - return WalkState::Continue; - }; - if !entry.file_type().is_some_and(|t| t.is_file()) { - return WalkState::Continue; - } - let Ok(md) = entry.metadata() else { - return WalkState::Continue; - }; - let is_config = entry - .file_name() - .to_str() - .is_some_and(crate::graph_batch::is_config_file); - if is_config { - if let Ok(rel) = entry.path().strip_prefix(root) { - let _ = tx.send(Seen::Config(config_line(&rel_string(rel), &md))); + let pruned: Pruned = Arc::default(); + walker_noting(dir, root, rules.clone(), Some(pruned.clone())) + .build_parallel() + .run(|| { + let tx = tx.clone(); + let rules = rules.clone(); + let cancel = cancel.clone(); + Box::new(move |entry| { + if cancel.is_cancelled() { + return WalkState::Quit; } - } - if let Some(found) = classify(root, entry.path(), &md, &rules) { - let _ = tx.send(Seen::File(found)); - } - WalkState::Continue - }) - }); + let Ok(entry) = entry else { + return WalkState::Continue; + }; + if !entry.file_type().is_some_and(|t| t.is_file()) { + return WalkState::Continue; + } + let Ok(md) = entry.metadata() else { + return WalkState::Continue; + }; + let is_config = entry + .file_name() + .to_str() + .is_some_and(crate::graph_batch::is_config_file); + if is_config { + if let Ok(rel) = entry.path().strip_prefix(root) { + let _ = tx.send(Seen::Config(config_line(&rel_string(rel), &md))); + } + } + if let Some(found) = classify(root, entry.path(), &md, &rules) { + let _ = tx.send(Seen::File(found)); + } + WalkState::Continue + }) + }); drop(tx); let mut kept = Vec::new(); - let mut skipped = Vec::new(); + let pruned = std::mem::take( + &mut *pruned + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner), + ); + let mut skipped = pruned_files(root, pruned, cancel); let mut configs = Vec::new(); for seen in rx { match seen { diff --git a/crates/atlas-codeindex/src/skip.rs b/crates/atlas-codeindex/src/skip.rs index c334f0f5b..44be6ac90 100644 --- a/crates/atlas-codeindex/src/skip.rs +++ b/crates/atlas-codeindex/src/skip.rs @@ -135,27 +135,57 @@ impl Rules { /// Why the file at `rel` (project-relative, `/`-separated) is skipped, or /// `None` to index it. + /// + /// A `!` line lifts only the built-in rule it negates: the skipped + /// directories at or above the path it re-includes (`!vendor/ours/` lifts + /// `vendor`), and the generated-name rule only when it names the file + /// itself. `node_modules`, `target`, generated files and the rest still + /// apply inside a re-included tree. pub(crate) fn path_verdict(&self, rel: &str) -> Option { - match self.atlasignore.matched_path_or_any_parents(rel, false) { + let lifted = match self.atlasignore.matched_path_or_any_parents(rel, false) { Match::Ignore(_) => return Some(SkipReason::AtlasIgnore), - Match::Whitelist(_) => return None, - Match::None => {} - } + Match::Whitelist(_) => self.reincluded_depth(rel), + Match::None => 0, + }; let mut parts = rel.split('/').collect::>(); let file = parts.pop().unwrap_or(rel); - if parts.iter().any(|d| is_skip_dir(d)) { + if parts + .iter() + .enumerate() + .any(|(i, d)| i >= lifted && is_skip_dir(d)) + { return Some(SkipReason::Vendor); } - if GENERATED_SUFFIXES.iter().any(|s| file.ends_with(s)) - || file.starts_with("zz_generated") - || file.contains(".generated.") + if lifted <= parts.len() + && (GENERATED_SUFFIXES.iter().any(|s| file.ends_with(s)) + || file.starts_with("zz_generated") + || file.contains(".generated.")) { return Some(SkipReason::GeneratedName); } None } - /// Whether the walk may skip the directory at `rel` entirely. + /// How many leading components of `rel` (a file) the deepest `!` line + /// matching it or one of its directories covers: 2 for `vendor/ours/a.rs` + /// under `!vendor/ours/`, all of them for a line naming the file. + fn reincluded_depth(&self, rel: &str) -> usize { + let parts: Vec<&str> = rel.split('/').collect(); + (1..=parts.len()) + .rev() + .find(|&k| { + let prefix = parts[..k].join("/"); + matches!( + self.atlasignore.matched(&prefix, k < parts.len()), + Match::Whitelist(_) + ) + }) + .unwrap_or(0) + } + + /// Whether the walk may skip the directory at `rel` entirely. A skipped + /// directory inside a re-included tree (`vendor/ours/node_modules` under + /// `!vendor/ours/`) is pruned: the re-include does not lift its rule. pub(crate) fn prune_dir(&self, rel: &str) -> bool { let name = rel.rsplit('/').next().unwrap_or(rel); let skipped = is_skip_dir(name) @@ -167,15 +197,67 @@ impl Rules { } } +/// Version-control directories: pruned, never counted as skipped code. +const VCS_DIRS: &[&str] = &[".git", ".hg", ".svn"]; + +impl Rules { + /// Why a directory [`Rules::prune_dir`] pruned was skipped, for the + /// skipped counts; `None` for version-control and Atlas's own directories, + /// which hold no project code. + pub(crate) fn prune_reason(&self, rel: &str) -> Option { + let name = rel.rsplit('/').next().unwrap_or(rel); + if VCS_DIRS.contains(&name) || atlas_search::is_atlas_dir(name) { + return None; + } + match self.atlasignore.matched_path_or_any_parents(rel, true) { + Match::Ignore(_) => Some(SkipReason::AtlasIgnore), + _ => Some(SkipReason::Vendor), + } + } +} + +/// Whether the walk descends into `name` when counting a pruned directory's +/// files: everything but version-control and Atlas's own directories. +pub(crate) fn counted_dir(name: &str) -> bool { + !VCS_DIRS.contains(&name) && !atlas_search::is_atlas_dir(name) +} + +/// Skipped files grouped as `(reason, directory, count)`, sorted by reason +/// then directory. A [`SkipReason::Vendor`] file counts under the directory +/// the built-in rule matched (`web/node_modules`); any other under its +/// top-level directory (`.` for a file at the root). +pub fn skipped_by_dir(skipped: &[(String, SkipReason)]) -> Vec<(String, String, usize)> { + let mut counts: std::collections::BTreeMap<(&str, String), usize> = Default::default(); + for (rel, reason) in skipped { + let parts: Vec<&str> = rel.split('/').collect(); + let dirs = &parts[..parts.len().saturating_sub(1)]; + let dir = match reason { + SkipReason::Vendor => dirs + .iter() + .position(|d| is_skip_dir(d)) + .map(|i| dirs[..=i].join("/")), + _ => None, + } + .unwrap_or_else(|| dirs.first().map_or(".", |d| d).to_string()); + *counts.entry((reason.label(), dir)).or_default() += 1; + } + counts + .into_iter() + .map(|((r, d), n)| (r.to_string(), d, n)) + .collect() +} + /// A built-in skipped directory name: [`SKIP_DIRS`] or Atlas's own. fn is_skip_dir(name: &str) -> bool { SKIP_DIRS.contains(&name) || atlas_search::is_atlas_dir(name) } -/// Whether re-include `pattern` (a `!` line without the `!`) can match a path -/// under directory `dir`. Gitignore rules: a pattern with no inner `/` -/// matches at any depth; otherwise it is anchored at the project root and -/// compared component by component (a glob component or `**` may match). +/// Whether re-include `pattern` (a `!` line without the `!`) can lift the skip +/// of directory `dir`: it names `dir` itself or something below it. Gitignore +/// rules: a pattern with no inner `/` matches at any depth; otherwise it is +/// anchored at the project root and compared component by component (a glob +/// component or `**` may match). A pattern that ends above `dir` re-includes +/// an ancestor, which leaves `dir`'s own skip in force. fn may_reach(pattern: &str, dir: &str) -> bool { let p = pattern.trim_end_matches('/'); if !p.contains('/') { @@ -184,8 +266,10 @@ fn may_reach(pattern: &str, dir: &str) -> bool { let mut parts = p.trim_start_matches('/').split('/'); for d in dir.split('/') { match parts.next() { - // `dir` lies inside what the pattern names, or `**` spans it. - None | Some("**") => return true, + // `dir` lies inside what the pattern names: its own skip stands. + None => return false, + // `**` spans it. + Some("**") => return true, Some(c) if c == d || c.contains(['*', '?', '[']) => {} Some(_) => return false, } diff --git a/crates/atlas-codeindex/src/store.rs b/crates/atlas-codeindex/src/store.rs index 6a328defc..00a969b80 100644 --- a/crates/atlas-codeindex/src/store.rs +++ b/crates/atlas-codeindex/src/store.rs @@ -16,7 +16,7 @@ use crate::IndexError; pub const SCHEMA_VERSION: i64 = 3; /// Bump when extraction output changes; a mismatch forces a full build. -pub const EXTRACTOR_VERSION: &str = "2.2"; +pub const EXTRACTOR_VERSION: &str = "2.3"; const SCHEMA_V1: &str = " CREATE TABLE meta(k TEXT PRIMARY KEY, v TEXT NOT NULL); diff --git a/crates/atlas-codeindex/src/tests/build.rs b/crates/atlas-codeindex/src/tests/build.rs index 7dc4ab080..66d8b4000 100644 --- a/crates/atlas-codeindex/src/tests/build.rs +++ b/crates/atlas-codeindex/src/tests/build.rs @@ -138,13 +138,16 @@ fn status_counts_skips_by_reason() { let ix = p.index(); let stats = ix.full_build(&CancelToken::new(), &|_| {}).unwrap(); assert_eq!(stats.files, 1); - assert_eq!(stats.skipped.len(), 4); + assert_eq!(stats.skipped.len(), 5); + // The pruned vendor/ counts too: a status that left it out would read + // as a complete index. assert_eq!( ix.status().unwrap().skipped, [ ("binary".to_string(), 1), ("generated_header".to_string(), 2), - ("generated_name".to_string(), 1) + ("generated_name".to_string(), 1), + ("vendor".to_string(), 1) ] ); } diff --git a/crates/atlas-codeindex/src/tests/graph.rs b/crates/atlas-codeindex/src/tests/graph.rs index 95243c2a0..48a4e9239 100644 --- a/crates/atlas-codeindex/src/tests/graph.rs +++ b/crates/atlas-codeindex/src/tests/graph.rs @@ -181,3 +181,246 @@ fn a_reconcile_sees_a_config_file_no_watcher_reported() { ix.reconcile(&atlas_search::CancelToken::new()).unwrap(); assert_eq!(importers(&ix), ["src/b.rs"]); } + +/// Call edges by caller and callee qualified name (`(src, dst)` pairs). +fn call_edges(ix: &CodeIndex) -> Vec<(String, String)> { + ix.with_reader(|c| { + let mut stmt = c.prepare( + "SELECT s.qualified_name, d.qualified_name FROM edges e + JOIN symbols s ON s.id = e.src_symbol_id JOIN symbols d ON d.id = e.dst_symbol_id + WHERE e.kind = 'call' ORDER BY 1, 2", + )?; + let rows = stmt + .query_map([], |r| Ok((r.get(0)?, r.get(1)?)))? + .collect::>>()?; + Ok(rows) + }) + .unwrap() +} + +fn calls(edges: &[(String, String)], src: &str, dst: &str) -> bool { + edges.iter().any(|(s, d)| s == src && d == dst) +} + +const CARGO_DEMO: &str = "[package]\nname = \"demo\"\nversion = \"0.1.0\"\nedition = \"2021\"\n"; + +/// A member call on a value whose type only a signature, a field or a smart pointer +/// names: rust-analyzer links all of these, and name matching alone linked none (the +/// method names are shared, or the receiver is untyped and the call crosses modules). +fn typed_receivers_project() -> (Project, CodeIndex) { + let p = Project::new(); + p.write("Cargo.toml", CARGO_DEMO); + p.write( + "src/lib.rs", + "pub mod store;\npub mod app;\npub mod other;\n", + ); + p.write( + "src/store.rs", + "use std::sync::{Arc, Mutex};\n\ + pub struct Item;\n\ + impl Item {\n pub fn name(&self) -> String { String::new() }\n pub fn len(&self) -> usize { 0 }\n}\n\ + pub struct Registry {\n items: Vec,\n}\n\ + impl Registry {\n pub fn open() -> std::io::Result { Ok(Registry { items: Vec::new() }) }\n pub fn find(&self, _id: u32) -> Option<&Item> { self.items.first() }\n pub fn size(&self) -> usize { self.items.len() }\n}\n\ + pub struct Store;\n\ + impl Store {\n pub fn request(&self) {}\n pub fn registry(&self) -> Arc { todo!() }\n}\n\ + pub fn new_store() -> (Store, u32) { (Store, 0) }\n\ + pub async fn connect() -> anyhow::Result { Ok(Store) }\n\ + pub struct Host {\n inner: Arc>,\n reg: Registry,\n}\n\ + impl Host {\n pub fn poke(&self) { self.inner.lock().unwrap().request(); }\n pub fn lookup(&self) -> String { self.reg.find(1).unwrap().name() }\n}\n", + ); + // Same method names elsewhere, so a name match is never unique. + p.write( + "src/other.rs", + "pub struct Decoy;\n\ + impl Decoy {\n pub fn request(&self) {}\n pub fn size(&self) -> usize { 0 }\n pub fn name(&self) -> String { String::new() }\n}\n", + ); + p.write( + "src/app.rs", + "use std::sync::Arc;\n\ + use crate::store::{connect, new_store, Registry, Store};\n\ + pub fn tuple_local() {\n let (store, _n) = new_store();\n store.request();\n}\n\ + pub async fn awaited() -> anyhow::Result<()> {\n let s = connect().await?;\n s.request();\n Ok(())\n}\n\ + pub fn opened() -> usize {\n let r = Registry::open().unwrap();\n r.size()\n}\n\ + pub fn param(s: Arc) {\n s.request();\n}\n\ + pub fn chained(s: &Store) -> usize {\n s.registry().size()\n}\n", + ); + let ix = p.built(); + (p, ix) +} + +#[test] +fn a_local_bound_from_a_call_is_typed_by_the_callees_return_type() { + let (_p, ix) = typed_receivers_project(); + let e = call_edges(&ix); + // `let (store, _) = new_store()`: the tuple's first element. + assert!(calls(&e, "tuple_local", "Store::request"), "{e:?}"); + // `connect().await?`: an async fn's `anyhow::Result`, peeled. + assert!(calls(&e, "awaited", "Store::request"), "{e:?}"); + // `Registry::open().unwrap()`: `io::Result`, with `Self` the impl type. + assert!(calls(&e, "opened", "Registry::size"), "{e:?}"); + for caller in ["tuple_local", "awaited", "opened"] { + assert!( + !e.iter() + .any(|(s, d)| s == caller && d.starts_with("Decoy::")), + "{e:?}" + ); + } +} + +#[test] +fn smart_pointers_fields_and_method_chains_type_a_receiver() { + let (_p, ix) = typed_receivers_project(); + let e = call_edges(&ix); + // A parameter `Arc` is a `Store` to a method call. + assert!(calls(&e, "param", "Store::request"), "{e:?}"); + // `s.registry()` returns `Arc`. + assert!(calls(&e, "chained", "Registry::size"), "{e:?}"); + // `self.inner: Arc>` through `.lock().unwrap()`. + assert!(calls(&e, "Host::poke", "Store::request"), "{e:?}"); + // `self.reg: Registry`, `find` returns `Option<&Item>`. + assert!(calls(&e, "Host::lookup", "Registry::find"), "{e:?}"); + assert!(calls(&e, "Host::lookup", "Item::name"), "{e:?}"); + // `self.items: Vec`: `len` is the Vec's, not `Item::len`. + assert!(!calls(&e, "Registry::size", "Item::len"), "{e:?}"); +} + +/// `tests/support/mod.rs` is declared by every test crate (`mod support;`), and each of +/// them calls into it; only the first crate to claim the file used to resolve. +#[test] +fn a_support_module_shared_by_test_crates_resolves_in_each() { + let p = Project::new(); + p.write("Cargo.toml", CARGO_DEMO); + // Same names in the library, so only the module path picks the right one. + p.write( + "src/lib.rs", + "pub fn helper() {}\npub mod fake {\n pub fn fake_server() {}\n}\n", + ); + p.write( + "tests/support/mod.rs", + "pub mod fake;\npub fn helper() {}\n", + ); + p.write("tests/support/fake.rs", "pub fn fake_server() {}\n"); + for t in ["a", "b"] { + p.write( + &format!("tests/{t}.rs"), + &format!( + "mod support;\nuse support::helper;\n#[test]\nfn t_{t}() {{\n helper();\n support::fake::fake_server();\n}}\n" + ), + ); + } + let ix = p.built(); + let e = call_edges(&ix); + for t in ["t_a", "t_b"] { + let to_support = |name: &str| { + ix.with_reader(|c| { + c.query_row( + "SELECT COUNT(*) FROM edges e JOIN symbols s ON s.id = e.src_symbol_id + JOIN symbols d ON d.id = e.dst_symbol_id JOIN files f ON f.id = d.file_id + WHERE s.name = ?1 AND d.name = ?2 AND f.rel LIKE 'tests/support/%'", + [t, name], + |r| r.get::<_, i64>(0), + ) + }) + .unwrap() + }; + assert_eq!(to_support("helper"), 1, "{t}: {e:?}"); + assert_eq!(to_support("fake_server"), 1, "{t}: {e:?}"); + } +} + +/// `use demo::*;` brings in a type whose method shares a name with a local `fn`: a bare +/// call is the local item (items shadow glob imports, and a bare name never calls a method). +#[test] +fn a_bare_call_prefers_a_local_fn_over_a_glob_imported_method() { + let p = Project::new(); + p.write("Cargo.toml", CARGO_DEMO); + p.write( + "src/lib.rs", + "pub struct Ev;\nimpl Ev {\n pub fn id(&self) -> u32 { 0 }\n}\npub fn shared() {}\n", + ); + p.write( + "tests/t.rs", + "use demo::*;\nfn id(x: u32) -> u32 { x }\n#[test]\nfn uses() {\n id(1);\n shared();\n}\n", + ); + let ix = p.built(); + let e = call_edges(&ix); + assert!(calls(&e, "uses", "id"), "{e:?}"); + assert!(!calls(&e, "uses", "Ev::id"), "{e:?}"); + assert!( + calls(&e, "uses", "shared"), + "the glob still resolves: {e:?}" + ); +} + +/// A typed chain reads another file's return types and fields: editing only that file +/// re-types the receiver (the calling file is unchanged, so only the surface rule reaches it). +#[test] +fn a_changed_return_type_or_field_re_types_a_receiver_in_another_file() { + let p = Project::new(); + p.write("Cargo.toml", CARGO_DEMO); + p.write( + "src/lib.rs", + "pub mod kinds;\npub mod make;\npub mod app;\n", + ); + p.write( + "src/kinds.rs", + "pub struct A;\nimpl A {\n pub fn go(&self) {}\n}\npub struct B;\nimpl B {\n pub fn go(&self) {}\n}\n", + ); + let make = |ret: &str, field: &str| { + format!( + "use crate::kinds::{{A, B}};\npub fn make() -> {ret} {{ todo!() }}\npub struct Holder {{\n pub inner: {field},\n}}\n" + ) + }; + p.write("src/make.rs", &make("A", "A")); + p.write( + "src/app.rs", + "use crate::make::{make, Holder};\npub fn run() {\n let x = make();\n x.go();\n}\npub fn held(h: &Holder) {\n h.inner.go();\n}\n", + ); + let ix = p.built(); + let e = call_edges(&ix); + assert!( + calls(&e, "run", "A::go") && calls(&e, "held", "A::go"), + "{e:?}" + ); + p.write("src/make.rs", &make("B", "B")); + ix.update_paths(&[p.path("src/make.rs")]).unwrap(); + let e = call_edges(&ix); + assert!( + calls(&e, "run", "B::go") && calls(&e, "held", "B::go"), + "{e:?}" + ); + assert!( + !calls(&e, "run", "A::go") && !calls(&e, "held", "A::go"), + "{e:?}" + ); +} + +/// `let (mut store, _) = new_store();` and `let c = Arc::new(Store::new());`. +#[test] +fn mutable_bindings_and_shared_pointers_keep_their_type() { + let p = Project::new(); + p.write("Cargo.toml", CARGO_DEMO); + p.write( + "src/lib.rs", + "pub mod store;\npub mod other;\npub mod app;\n", + ); + p.write( + "src/store.rs", + "pub struct Store;\nimpl Store {\n pub fn new() -> Self { Store }\n pub fn request(&mut self) {}\n}\npub fn new_store() -> (Store, u32) { (Store, 0) }\n", + ); + p.write( + "src/other.rs", + "pub struct Decoy;\nimpl Decoy {\n pub fn request(&self) {}\n}\n", + ); + p.write( + "src/app.rs", + "use std::sync::Arc;\nuse crate::store::{new_store, Store};\n\ + pub fn tuple_mut() {\n let (mut store, _n) = new_store();\n store.request();\n}\n\ + pub fn shared() {\n let c = Arc::new(Store::new());\n c.request();\n}\n", + ); + let ix = p.built(); + let e = call_edges(&ix); + assert!(calls(&e, "tuple_mut", "Store::request"), "{e:?}"); + assert!(calls(&e, "shared", "Store::request"), "{e:?}"); + assert!(!e.iter().any(|(_, d)| d == "Decoy::request"), "{e:?}"); +} diff --git a/crates/atlas-codeindex/src/tests/query.rs b/crates/atlas-codeindex/src/tests/query.rs index 09ce42ef2..b4e18d048 100644 --- a/crates/atlas-codeindex/src/tests/query.rs +++ b/crates/atlas-codeindex/src/tests/query.rs @@ -57,6 +57,31 @@ fn camel_and_snake_parts_match_either_spelling() { ); } +#[test] +fn exact_count_says_whether_any_hit_is_the_name_asked_for() { + let p = search_fixture(); + let ix = p.built(); + // Exact name, case-blind name and qualified suffix count, on every page. + for (query, offset) in [ + ("Store", 0), + ("Store", 3), + ("storeview", 0), + ("Store::open", 0), + ] { + let (_, _, exact) = ix + .find_symbol_with_exact(&SymbolQuery { offset, ..q(query) }) + .unwrap(); + assert!(exact > 0, "{query} @ {offset}"); + } + // A prefix or word match is not the symbol asked for, however many hits. + for query in ["Stor", "cloud client", "StoreFrobnicator"] { + let (hits, _, exact) = ix.find_symbol_with_exact(&q(query)).unwrap(); + assert_eq!(exact, 0, "{query}: {:?}", names(&hits)); + } + let (hits, _, _) = ix.find_symbol_with_exact(&q("Stor")).unwrap(); + assert!(!hits.is_empty()); +} + #[test] fn qualified_name_query_hits_the_method() { let p = search_fixture(); diff --git a/crates/atlas-codeindex/src/tests/scan.rs b/crates/atlas-codeindex/src/tests/scan.rs index 2572f4617..f04738d8f 100644 --- a/crates/atlas-codeindex/src/tests/scan.rs +++ b/crates/atlas-codeindex/src/tests/scan.rs @@ -41,10 +41,14 @@ fn discover_sorts_prunes_vendor_and_honours_gitignore() { .write("README.md", "# not code\n"); let (kept, skipped) = walk(&p); assert_eq!(rels(&kept), ["src/a.rs", "web/b.ts"]); - // vendor/ is pruned during the walk, so it is never even listed. + // vendor/ is pruned during the walk, but its source files are still + // counted as skipped: an index missing a directory must say so. assert_eq!( skipped, - [("api/api.pb.go".to_string(), SkipReason::GeneratedName)] + [ + ("api/api.pb.go".to_string(), SkipReason::GeneratedName), + ("vendor/lib/x.go".to_string(), SkipReason::Vendor), + ] ); } @@ -58,10 +62,14 @@ fn atlasignore_reinclude_walks_into_vendor() { .write("src/main.rs", "fn main() {}\n"); let (kept, skipped) = walk(&p); assert_eq!(rels(&kept), ["src/main.rs", "vendor/ours/a.rs"]); - // `experiments/` is pruned unwalked: no re-include reaches into it. + // `experiments/` is pruned (no re-include reaches into it), and still + // reported. assert_eq!( skipped, - [("vendor/theirs/b.rs".to_string(), SkipReason::Vendor)] + [ + ("experiments/x.rs".to_string(), SkipReason::AtlasIgnore), + ("vendor/theirs/b.rs".to_string(), SkipReason::Vendor), + ] ); } @@ -143,3 +151,38 @@ fn process_all_keeps_path_order_and_honours_known_hashes() { .iter() .all(|o| matches!(o, Outcome::Cancelled))); } + +#[test] +fn pruned_dirs_report_their_source_files_but_not_vcs_or_gitignored_ones() { + let p = Project::new(); + p.write("src/a.rs", "pub fn a() {}\n") + .write("web/node_modules/pkg/index.js", "module.exports = 1;\n") + .write("web/node_modules/pkg/README.md", "# not code\n") + .write("third_party/deep/vendor/x.go", "package x\n") + .write(".gitignore", "target/\n") + .write("target/debug/build.rs", "fn main() {}\n") + .write(".git/hooks/pre_commit.py", "x = 0\n"); + let (kept, skipped) = walk(&p); + assert_eq!(rels(&kept), ["src/a.rs"]); + // Gitignored dirs were never the index's to skip; `.git` is not code. + assert_eq!( + skipped, + [ + ( + "third_party/deep/vendor/x.go".to_string(), + SkipReason::Vendor + ), + ( + "web/node_modules/pkg/index.js".to_string(), + SkipReason::Vendor + ), + ] + ); + assert_eq!( + crate::skip::skipped_by_dir(&skipped), + [ + ("vendor".to_string(), "third_party".to_string(), 1), + ("vendor".to_string(), "web/node_modules".to_string(), 1), + ] + ); +} diff --git a/crates/atlas-codeindex/src/tests/skip.rs b/crates/atlas-codeindex/src/tests/skip.rs index 4349f71ac..ca5566464 100644 --- a/crates/atlas-codeindex/src/tests/skip.rs +++ b/crates/atlas-codeindex/src/tests/skip.rs @@ -68,6 +68,40 @@ fn atlasignore_excludes_and_reincludes() { ); } +/// A `!` line lifts the rule it negates and nothing else: re-including +/// `vendor/ours/` indexes that tree despite the `vendor` skip, but the +/// built-in skips still apply inside it. +#[test] +fn a_reinclude_lifts_only_the_skip_it_names() { + let p = Project::new(); + p.write( + ".atlasignore", + "!vendor/ours/\n!keep/node_modules/\n!src/schema.generated.ts\n", + ); + let rules = Rules::load(p.root()); + assert_eq!(rules.path_verdict("vendor/ours/src/lib.rs"), None); + for (rel, why) in [ + ("vendor/ours/node_modules/x/index.js", SkipReason::Vendor), + ("vendor/ours/target/debug/build.rs", SkipReason::Vendor), + ("vendor/ours/web/dist/app.js", SkipReason::Vendor), + ("vendor/ours/api/v1/api.pb.go", SkipReason::GeneratedName), + ("vendor/ours/web/app.min.js", SkipReason::GeneratedName), + ] { + assert_eq!(rules.path_verdict(rel), Some(why), "{rel}"); + } + // The skipped directories inside a re-included tree are pruned again. + assert!(!rules.prune_dir("vendor") && !rules.prune_dir("vendor/ours")); + assert!(rules.prune_dir("vendor/ours/node_modules") && rules.prune_dir("vendor/ours/target")); + // A re-include naming a skipped directory itself, or a file, lifts that. + assert_eq!(rules.path_verdict("keep/node_modules/a.js"), None); + assert!(!rules.prune_dir("keep/node_modules")); + assert_eq!(rules.path_verdict("src/schema.generated.ts"), None); + assert_eq!( + rules.path_verdict("src/other.generated.ts"), + Some(SkipReason::GeneratedName) + ); +} + #[test] fn ignore_chain_follows_nested_gitignores_and_info_exclude() { let p = Project::new(); @@ -103,3 +137,26 @@ fn a_project_inside_a_repository_honours_ignores_above_it() { ); assert!(!chain.is_ignored("src/main.ts", false)); } + +/// This repository's own `.atlasignore`: the engine under `vendor/` is +/// first-party code (ADR-0003) and must be indexed, while the generic +/// `vendor` skip still applies to anything else vendored. +#[test] +fn this_repository_indexes_its_vendored_engine() { + let root = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../.."); + let rules = Rules::load(&root); + assert_eq!( + rules.path_verdict("vendor/atlas-engine/app-server-client/src/lib.rs"), + None + ); + assert!(!rules.prune_dir("vendor") && !rules.prune_dir("vendor/atlas-engine")); + assert_eq!( + rules.path_verdict("vendor/other/x.rs"), + Some(SkipReason::Vendor) + ); + assert_eq!( + rules.path_verdict("vendor/atlas-engine/x/node_modules/y.js"), + Some(SkipReason::Vendor) + ); + assert!(rules.prune_dir("vendor/atlas-engine/target")); +} diff --git a/crates/atlas-codeindex/src/universe.rs b/crates/atlas-codeindex/src/universe.rs index 46c419dd1..1b44a4692 100644 --- a/crates/atlas-codeindex/src/universe.rs +++ b/crates/atlas-codeindex/src/universe.rs @@ -42,6 +42,9 @@ pub(crate) struct SymRow { pub qn: String, pub exported: bool, pub is_test: bool, + /// Header text (`pub fn open() -> io::Result`): Rust receiver typing reads return + /// types and type aliases from it. + pub sig: String, } #[derive(Debug, Clone)] @@ -69,6 +72,8 @@ pub(crate) struct Universe { pub by_file_name: HashMap<(i64, String), Vec>, /// (family, module) → file ids, ascending. pub by_module: HashMap<(u8, String), Vec>, + /// Rust struct symbol id → its fields `(name, type text)`, in declaration order. + pub fields: HashMap>, } impl Universe { @@ -85,7 +90,7 @@ impl Universe { })? .collect::>>()?; let syms = conn - .prepare("SELECT id, file_id, parent_id, kind, name, qualified_name, exported, is_test FROM symbols ORDER BY id")? + .prepare("SELECT id, file_id, parent_id, kind, name, qualified_name, exported, is_test, signature FROM symbols ORDER BY id")? .query_map([], |r| { Ok(SymRow { id: r.get(0)?, @@ -96,6 +101,7 @@ impl Universe { qn: r.get(5)?, exported: r.get::<_, i64>(6)? != 0, is_test: r.get::<_, i64>(7)? != 0, + sig: r.get(8)?, }) })? .collect::>>()?; @@ -133,7 +139,23 @@ impl Universe { let (fid, m) = row?; mods.entry(fid).or_default().push(m); } - Ok(Universe::from_parts(files, syms, imports, mods)) + let mut u = Universe::from_parts(files, syms, imports, mods); + let mut stmt = conn.prepare( + "SELECT src_symbol_id, name, receiver FROM refs + WHERE kind = 'field' AND src_symbol_id IS NOT NULL ORDER BY id", + )?; + let rows = stmt.query_map([], |r| { + Ok(( + r.get::<_, i64>(0)?, + r.get::<_, String>(1)?, + r.get::<_, String>(2)?, + )) + })?; + for row in rows { + let (sid, name, ty) = row?; + u.fields.entry(sid).or_default().push((name, ty)); + } + Ok(u) } pub(crate) fn from_parts( diff --git a/scripts/ci-affected.mjs b/scripts/ci-affected.mjs index ce51fdb31..ca84f33e6 100644 --- a/scripts/ci-affected.mjs +++ b/scripts/ci-affected.mjs @@ -109,6 +109,12 @@ export const EXTRA_INPUTS = [ testOnly: true, why: "the tests module of src/lib.rs, and tests/import.rs, check the mock's snapshots", }, + { + path: ".atlasignore", + packages: ["atlas-codeindex"], + testOnly: true, + why: "src/tests/skip.rs checks this repository's own .atlasignore indexes the engine", + }, { path: "docs/agents/delta-wire-contract.md", packages: ["atlas-agent-wire"], @@ -157,6 +163,8 @@ export const NON_RUST = [ ".oxfmtrc.json", ".gitignore", ".env.example", + // Also an EXTRA_INPUTS entry, which still marks atlas-codeindex. + ".atlasignore", ]; /** diff --git a/src-tauri/src/commands/code_index/symbols.rs b/src-tauri/src/commands/code_index/symbols.rs index bde5d0ab8..c7c04ee99 100644 --- a/src-tauri/src/commands/code_index/symbols.rs +++ b/src-tauri/src/commands/code_index/symbols.rs @@ -247,16 +247,22 @@ fn find_symbol(scope: &Scope, args: &Value) -> Result { limit: arg_usize(args, "limit").unwrap_or(20).min(MAX_FIND_LIMIT), offset: arg_usize(args, "offset").unwrap_or(0), }; - let (hits, page) = scope + let (hits, page, exact) = scope .project .index - .find_symbol(&q) + .find_symbol_with_exact(&q) .map_err(|e| e.to_string())?; if hits.is_empty() { return Ok(format!( "symbols: 0 for {query:?}. Try fewer words, another spelling, or grep." )); } + // Word and prefix matches look like an answer; say when none is the name. + let header = if exact == 0 { + format!("no exact match for {query:?}; closest by words:\n") + } else { + String::new() + }; let rows = hits .iter() .map(|h| { @@ -273,7 +279,8 @@ fn find_symbol(scope: &Scope, args: &Value) -> Result { cols: vec!["symbol", "kind", "loc", "signature"], rows, }; - Ok(render(&[table], &page, budget(args))) + let budget = budget(args).saturating_sub(header.len()); + Ok(header + &render(&[table], &page, budget)) } /// One row per symbol, names indented under their parent. @@ -559,6 +566,29 @@ mod tests { assert!(outline.contains("outside the session root"), "{outline}"); } + #[test] + fn find_symbol_says_when_nothing_is_the_name_asked_for() { + let dir = project(); + let reg = ready(dir.path()); + let scope = Scope::resolve(®, dir.path()).unwrap(); + let found = call(&scope, "find_symbol", &json!({ "query": "openStor" })).unwrap(); + assert!( + found.starts_with("no exact match for \"openStor\"; closest by words:\nsymbols: "), + "{found}" + ); + assert!(found.contains("web/app.ts:1-1"), "{found}"); + // A real name gets the plain table, even on a later page. + for offset in [0, 1] { + let found = call( + &scope, + "find_symbol", + &json!({ "query": "Store", "offset": offset }), + ) + .unwrap(); + assert!(!found.contains("no exact match"), "{found}"); + } + } + #[test] fn long_symbol_lists_members_and_unknown_name_suggests() { let dir = project(); From 42833f5e99888b16c560e905570b4b4222f60821 Mon Sep 17 00:00:00 2001 From: Uzayer Masud <16221196+uzayer@users.noreply.github.com> Date: Sat, 10 Oct 2026 00:00:31 +0600 Subject: [PATCH 2/3] fix(models): an empty model folder is not an install, and the composer offers the code embedding model - looking up a model's path no longer creates its folder; installed means every model file is present, and model_download creates the folder on tokio::fs - an 'Enable semantic search' pill in the composer starts the download or opens Local Models; hidden while agent code tools are off, dismissible per model (every open composer hears it), never downloads on its own - the mock backend keeps memory and code model selections apart and fakes a download Co-authored-by: Ahammad Nafiz --- src-tauri/src/commands/models.rs | 109 +++++++++++-- src/components/code-model-nudge.test.tsx | 151 ++++++++++++++++++ src/components/code-model-nudge.tsx | 100 ++++++++++++ src/dev/mock-backend/fixtures/misc.ts | 57 ++++++- .../chat/components/message-input.tsx | 4 + src/features/settings/lib/code-model-nudge.ts | 51 ++++++ src/features/settings/stores/models-store.ts | 9 ++ 7 files changed, 462 insertions(+), 19 deletions(-) create mode 100644 src/components/code-model-nudge.test.tsx create mode 100644 src/components/code-model-nudge.tsx create mode 100644 src/features/settings/lib/code-model-nudge.ts diff --git a/src-tauri/src/commands/models.rs b/src-tauri/src/commands/models.rs index 83c1bc4b7..3feefbb80 100644 --- a/src-tauri/src/commands/models.rs +++ b/src-tauri/src/commands/models.rs @@ -202,22 +202,25 @@ pub fn find_entry(id: &str) -> Option { // ── Selection + path resolution (the chokepoints delegate here) ──────────────── -/// Root dir holding every downloaded model: `app_data/models/`. +/// Root dir holding every downloaded model: `app_data/models/`. A pure path +/// lookup: nothing is created until a download needs the directory. pub fn models_root(app: &AppHandle) -> Result { - let dir = app + Ok(app .path() .app_data_dir() .map_err(|e| format!("no app data dir: {e}"))? - .join("models"); - std::fs::create_dir_all(&dir).map_err(|e| format!("create models dir: {e}"))?; - Ok(dir) + .join("models")) } -/// On-disk dir for a model id (created lazily). +/// On-disk dir for a model id. A pure path lookup, like [`models_root`]: an +/// empty directory left behind by a lookup would read as a broken install, so +/// only a download ([`model_download`], [`download_files`]) creates it. pub fn model_dir_for(app: &AppHandle, id: &str) -> Result { - let dir = models_root(app)?.join(id); - std::fs::create_dir_all(&dir).map_err(|e| format!("create model dir: {e}"))?; - Ok(dir) + Ok(model_dir_in(&models_root(app)?, id)) +} + +fn model_dir_in(root: &Path, id: &str) -> PathBuf { + root.join(id) } /// The memory model to use. A code-only model is never used for memory, @@ -243,13 +246,16 @@ pub(crate) fn code_model_spec(id: &str) -> Option { /// Whether every file for `id` exists on disk. Uses the catalog's file list; for an /// unknown id (e.g. a legacy/manual dir) falls back to the embedding file triplet. +/// The directory alone, or the spec written ahead of the weights, is not an install. pub fn is_downloaded(app: &AppHandle, id: &str) -> bool { - let Ok(dir) = model_dir_for(app, id) else { - return false; - }; + model_dir_for(app, id).is_ok_and(|dir| files_present(&dir, id)) +} + +fn files_present(dir: &Path, id: &str) -> bool { + let present = |f: &str| dir.join(f).is_file(); match find_entry(id) { - Some(e) => e.dest_files().iter().all(|f| dir.join(f).exists()), - None => EMBED_FILES.iter().all(|f| dir.join(f).exists()), + Some(e) => e.dest_files().into_iter().all(present), + None => EMBED_FILES.into_iter().all(present), } } @@ -280,6 +286,11 @@ pub async fn download_files( use futures::StreamExt; use tokio::io::AsyncWriteExt; + // The one place a model's directory comes into being (plus `model_download`, + // which writes the spec first). + tokio::fs::create_dir_all(dir) + .await + .map_err(|e| format!("create model dir: {e}"))?; let client = reqwest::Client::builder() .user_agent("Atlas-IDE") .build() @@ -402,6 +413,9 @@ pub async fn models_list(app: AppHandle) -> Result, String> { pub async fn model_download(app: AppHandle, id: String) -> Result<(), String> { let entry = find_entry(&id).ok_or_else(|| format!("unknown model '{id}'"))?; let dir = model_dir_for(&app, &id)?; + tokio::fs::create_dir_all(&dir) + .await + .map_err(|e| format!("create model dir: {e}"))?; // A code model's spec goes next to its weights before the download starts // (it is not one of `EMBED_FILES`, so it never makes a model count as // downloaded on its own). @@ -446,7 +460,12 @@ pub async fn model_remove(app: AppHandle, id: String) -> Result<(), String> { return Err("Can't remove the model that's currently selected.".into()); } let dir = model_dir_for(&app, &id)?; - std::fs::remove_dir_all(&dir).map_err(|e| format!("remove {id}: {e}"))?; + match std::fs::remove_dir_all(&dir) { + Ok(()) => {} + // Nothing on disk is already the state removal asks for. + Err(e) if e.kind() == std::io::ErrorKind::NotFound => {} + Err(e) => return Err(format!("remove {id}: {e}")), + } let _ = app.emit("atlas:models-changed", ()); Ok(()) } @@ -578,6 +597,66 @@ mod tests { assert_eq!(memory_model_id("my-local-model".into()), "my-local-model"); } + #[test] + fn looking_up_a_model_dir_creates_nothing() { + let tmp = tempfile::tempdir().unwrap(); + let root = tmp.path().join("models"); + let dir = model_dir_in(&root, "granite-embedding-small-r2"); + assert_eq!(dir, root.join("granite-embedding-small-r2")); + assert!( + !root.exists(), + "a path lookup must not create the models root" + ); + // And asking whether it is installed must not create it either. + assert!(!files_present(&dir, "granite-embedding-small-r2")); + assert!( + !dir.exists(), + "an installed check must not create the model dir" + ); + } + + #[test] + fn an_empty_or_partial_model_dir_is_not_installed() { + let tmp = tempfile::tempdir().unwrap(); + let id = "granite-embedding-small-r2"; + let dir = model_dir_in(tmp.path(), id); + std::fs::create_dir_all(&dir).unwrap(); + assert!(!files_present(&dir, id), "an empty dir is not an install"); + + // The spec is written before the weights arrive; on its own it is not an install. + code_model_spec(id).unwrap().write(&dir).unwrap(); + std::fs::write(dir.join("config.json"), "{}").unwrap(); + std::fs::write(dir.join("model.safetensors.part"), "half").unwrap(); + assert!( + !files_present(&dir, id), + "a half-finished download is not an install" + ); + + // A directory where a file belongs is not the file. + std::fs::create_dir(dir.join("tokenizer.json")).unwrap(); + std::fs::write(dir.join("model.safetensors"), "weights").unwrap(); + assert!(!files_present(&dir, id)); + + std::fs::remove_dir(dir.join("tokenizer.json")).unwrap(); + std::fs::write(dir.join("tokenizer.json"), "{}").unwrap(); + assert!( + files_present(&dir, id), + "every model file present is an install" + ); + } + + #[test] + fn an_unknown_model_dir_needs_the_embedding_triplet() { + let tmp = tempfile::tempdir().unwrap(); + let dir = model_dir_in(tmp.path(), "my-local-model"); + std::fs::create_dir_all(&dir).unwrap(); + assert!(!files_present(&dir, "my-local-model")); + for f in EMBED_FILES { + std::fs::write(dir.join(f), "x").unwrap(); + } + assert!(files_present(&dir, "my-local-model")); + } + #[test] fn the_wire_name_is_snake_case() { assert_eq!( diff --git a/src/components/code-model-nudge.test.tsx b/src/components/code-model-nudge.test.tsx new file mode 100644 index 000000000..a998f8120 --- /dev/null +++ b/src/components/code-model-nudge.test.tsx @@ -0,0 +1,151 @@ +// @vitest-environment happy-dom +import { act, cleanup, fireEvent, render, screen } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import type { ModelStatus } from "@/features/settings/lib/models-api"; +import { + codeModelToSuggest, + downloadPercent, + readDismissed, +} from "@/features/settings/lib/code-model-nudge"; +import { DEFAULT_SETTINGS } from "@/features/settings/lib/app-settings"; + +const invoke = vi.fn(); +vi.mock("@tauri-apps/api/core", () => ({ invoke: (...a: unknown[]) => invoke(...a) })); +vi.mock("@tauri-apps/api/event", () => ({ listen: vi.fn(async () => () => {}) })); +const openSettingsSection = vi.fn(); +vi.mock("@/features/settings/lib/open-settings", () => ({ + openSettingsSection: (s: string) => openSettingsSection(s), +})); + +const { CodeModelNudge } = await import("./code-model-nudge"); +const { useModelsStore } = await import("@/features/settings/stores/models-store"); +const { useSettingsStore } = await import("@/features/settings/stores/settings-store"); + +function model(over: Partial): ModelStatus { + return { + id: "granite-embedding-small-r2", + kind: "code_embedding", + name: "Granite Embedding Small R2", + repo: "ibm-granite/granite-embedding-small-english-r2", + files: [], + dim: 384, + sizeMb: 95, + description: "", + compatible: true, + downloaded: false, + selected: true, + ...over, + }; +} + +const memory = model({ id: "all-MiniLM-L6-v2", kind: "embedding", downloaded: false }); + +beforeEach(() => { + cleanup(); + localStorage.clear(); + invoke.mockReset(); + openSettingsSection.mockReset(); + useModelsStore.setState({ + list: [], + loaded: false, + downloading: {}, + pending: null, + codeNudgeDismissed: [], + }); + useSettingsStore.setState({ settings: { ...DEFAULT_SETTINGS, agentCodeTools: true } }); +}); + +describe("codeModelToSuggest", () => { + it("suggests the selected code model while it is missing", () => { + expect(codeModelToSuggest([memory, model({})], [])?.id).toBe("granite-embedding-small-r2"); + }); + + it("is quiet once it is downloaded, dismissed, or when only memory's model is missing", () => { + expect(codeModelToSuggest([model({ downloaded: true })], [])).toBeNull(); + expect(codeModelToSuggest([model({})], ["granite-embedding-small-r2"])).toBeNull(); + expect(codeModelToSuggest([memory, model({ selected: false })], [])).toBeNull(); + expect(codeModelToSuggest([], [])).toBeNull(); + }); + + it("asks again for a different code model", () => { + expect( + codeModelToSuggest([model({ id: "coderankembed" })], ["granite-embedding-small-r2"])?.id, + ).toBe("coderankembed"); + }); +}); + +describe("downloadPercent", () => { + it("counts finished files and the one in flight", () => { + expect(downloadPercent({ fileIndex: 0, fileCount: 1, received: 0, total: 0 })).toBe(0); + expect(downloadPercent({ fileIndex: 2, fileCount: 3, received: 50, total: 100 })).toBe(83); + expect(downloadPercent({ fileIndex: 2, fileCount: 3, received: 100, total: 100 })).toBe(100); + }); +}); + +describe("CodeModelNudge", () => { + const list = [memory, model({})]; + + function renderWithCatalog() { + invoke.mockImplementation(async (cmd: string) => (cmd === "models_list" ? list : null)); + render(); + } + + it("one click starts the download, and never on its own", async () => { + renderWithCatalog(); + const button = await screen.findByRole("button", { name: "Enable semantic search" }); + expect(invoke).not.toHaveBeenCalledWith("model_download", expect.anything()); + + await act(async () => fireEvent.click(button)); + expect(invoke).toHaveBeenCalledWith("model_download", { id: "granite-embedding-small-r2" }); + expect(await screen.findByRole("button", { name: "Downloading 0%" })).toBeTruthy(); + }); + + it("opens Local Models", async () => { + renderWithCatalog(); + fireEvent.click(await screen.findByRole("button", { name: "Open Local Models" })); + expect(openSettingsSection).toHaveBeenCalledWith("models"); + }); + + it("dismissing hides it and remembers that", async () => { + renderWithCatalog(); + fireEvent.click(await screen.findByRole("button", { name: "Dismiss" })); + expect(screen.queryByRole("group", { name: "Semantic code search" })).toBeNull(); + expect(readDismissed()).toEqual(["granite-embedding-small-r2"]); + }); + + it("dismissing in one composer hides it in every other open one", async () => { + invoke.mockImplementation(async (cmd: string) => (cmd === "models_list" ? list : null)); + render( + <> + + + , + ); + const [first] = await screen.findAllByRole("button", { name: "Dismiss" }); + expect(screen.getAllByRole("group", { name: "Semantic code search" })).toHaveLength(2); + fireEvent.click(first); + expect(screen.queryAllByRole("group", { name: "Semantic code search" })).toHaveLength(0); + }); + + it("stays hidden while agents have no code tools, and appears when they are turned on", async () => { + useSettingsStore.setState({ settings: { ...DEFAULT_SETTINGS, agentCodeTools: false } }); + renderWithCatalog(); + await act(async () => {}); + expect(screen.queryByRole("group", { name: "Semantic code search" })).toBeNull(); + + act(() => + useSettingsStore.setState({ settings: { ...DEFAULT_SETTINGS, agentCodeTools: true } }), + ); + expect(await screen.findByRole("group", { name: "Semantic code search" })).toBeTruthy(); + }); + + it("renders nothing once the model is on disk", async () => { + invoke.mockImplementation(async (cmd: string) => + cmd === "models_list" ? [model({ downloaded: true })] : null, + ); + render(); + await act(async () => {}); + expect(screen.queryByRole("group", { name: "Semantic code search" })).toBeNull(); + }); +}); diff --git a/src/components/code-model-nudge.tsx b/src/components/code-model-nudge.tsx new file mode 100644 index 000000000..29ea0de40 --- /dev/null +++ b/src/components/code-model-nudge.tsx @@ -0,0 +1,100 @@ +import { useEffect } from "react"; +import { toast } from "sonner"; +import { Boxes, Download, Loader2, X } from "lucide-react"; +import { cn } from "@/lib/utils"; +import { Hint } from "@/ui/tooltip"; +import { useModelsStore } from "@/features/settings/stores/models-store"; +import { useSettingsStore } from "@/features/settings/stores/settings-store"; +import { openSettingsSection } from "@/features/settings/lib/open-settings"; +import { codeModelToSuggest, downloadPercent } from "@/features/settings/lib/code-model-nudge"; + +const PART = + "flex items-center h-full text-muted-foreground hover:bg-element-hover hover:text-foreground transition-colors cursor-pointer"; + +/** + * A composer pill that offers the code search model while it is missing: + * one click downloads it, the gear opens Settings → Local Models, the cross + * hides it for good (per model, in every composer at once). Never downloads + * on its own: the model is large enough that it should be the user's call. + * Hidden while agents get no code tools (`agentCodeTools`, ADR-0015): with + * no `semantic_search` to serve, the model would buy nothing. + * + * Lives in `src/components` because it is a cross-feature widget: the + * settings feature owns its state, the chat composer renders it. + */ +export function CodeModelNudge({ labelClassName }: { labelClassName?: string }) { + const list = useModelsStore.use.list(); + const downloading = useModelsStore.use.downloading(); + const actions = useModelsStore.use.actions(); + const dismissed = useModelsStore.use.codeNudgeDismissed(); + const codeTools = useSettingsStore((s) => s.settings.agentCodeTools); + + useEffect(() => { + void actions.init(); + }, [actions]); + + const model = codeTools ? codeModelToSuggest(list, dismissed) : null; + if (!model) return null; + + const progress = downloading[model.id]; + const label = progress ? `Downloading ${downloadPercent(progress)}%` : "Enable semantic search"; + const hint = progress + ? `Downloading ${model.name} for semantic code search` + : `Code search is matching keywords only. Download ${model.name} (${model.sizeMb} MB) so plain-English questions find code too.`; + + const download = () => { + if (progress) return; + actions.download(model.id).catch((e) => + toast.error(`Download failed: ${e instanceof Error ? e.message : String(e)}`, { + action: { label: "Local Models", onClick: () => openSettingsSection("models") }, + }), + ); + }; + + return ( +
+ + + + + + + {!progress && ( + + + + )} +
+ ); +} diff --git a/src/dev/mock-backend/fixtures/misc.ts b/src/dev/mock-backend/fixtures/misc.ts index 13df55192..6acbb5b8a 100644 --- a/src/dev/mock-backend/fixtures/misc.ts +++ b/src/dev/mock-backend/fixtures/misc.ts @@ -261,10 +261,57 @@ const MODELS: ModelStatus[] = [ downloaded: false, selected: false, }, + // Code search's model: selected (it is the default) but not downloaded, the + // state a fresh install is in, so the composer's "Enable semantic search" + // pill shows. Downloading it walks the progress events below. + { + id: "granite-embedding-small-r2", + kind: "code_embedding", + name: "Granite Embedding Small R2", + repo: "ibm-granite/granite-embedding-small-english-r2", + files: ["config.json", "tokenizer.json", "model.safetensors"].map((file) => ({ + repo: "ibm-granite/granite-embedding-small-english-r2", + file, + dest: file, + })), + dim: 384, + sizeMb: 95, + description: + "Small, fast code search embeddings (47M, Apache-2.0). Default for semantic code search.", + compatible: true, + downloaded: false, + selected: true, + }, ]; const models: ModelStatus[] = MODELS.map((model) => ({ ...model })); +/** A download that takes a few seconds, so the progress states are visible. */ +function simulateDownload(model: ModelStatus): void { + const fileCount = model.files.length; + const total = model.sizeMb * 1_000_000; + const steps = 12; + for (let step = 1; step <= steps; step++) { + setTimeout(() => { + const done = step / steps; + const fileIndex = Math.min(fileCount - 1, Math.floor(done * fileCount)); + void emit("atlas:model-download:progress", { + id: model.id, + file: model.files[fileIndex]?.dest ?? "", + fileIndex, + fileCount, + received: Math.round(total * (done * fileCount - fileIndex)), + total, + }); + if (step === steps) { + model.downloaded = true; + void emit("atlas:model-download:done", { id: model.id, success: true, error: null }); + void emit("atlas:models-changed"); + } + }, step * 300); + } +} + // ── handlers ────────────────────────────────────────────────────────────── export const miscHandlers: MockHandlers = { @@ -372,8 +419,7 @@ export const miscHandlers: MockHandlers = { const model = models.find((candidate) => candidate.id === String(id)); if (!model) throw new Error(`unknown model ${String(id)}`); if (!model.compatible) throw new Error("not enough memory for this model"); - model.downloaded = true; - void emit("atlas:models-changed"); + simulateDownload(model); return null; }, model_remove: ({ id }): null => { @@ -388,8 +434,11 @@ export const miscHandlers: MockHandlers = { model_select: ({ id }): SelectResult => { const next = models.find((candidate) => candidate.id === String(id)); if (!next) throw new Error(`unknown model ${String(id)}`); - const previous = models.find((candidate) => candidate.selected); - for (const model of models) model.selected = model.id === next.id; + const previous = models.find((candidate) => candidate.selected && candidate.kind === next.kind); + // Each kind has its own selection, as in Rust. + for (const model of models) { + if (model.kind === next.kind) model.selected = model.id === next.id; + } void emit("atlas:models-changed"); // Changing dimension is what forces the re-index prompt. return { needsReindex: previous?.dim !== next.dim }; diff --git a/src/features/chat/components/message-input.tsx b/src/features/chat/components/message-input.tsx index 3ac981a8c..46baef902 100644 --- a/src/features/chat/components/message-input.tsx +++ b/src/features/chat/components/message-input.tsx @@ -76,6 +76,7 @@ import type { import { commandRequiresArgs } from "./slash-command-picker"; import { PlanTasksPill } from "./plan-tasks-pill"; import { openSettingsSection } from "@/features/settings/lib/open-settings"; +import { CodeModelNudge } from "@/components/code-model-nudge"; import { ComposerOptionsPill } from "./composer-options-pill"; import { UsagePill } from "./usage-pill"; import { composerPillLabelClass } from "./composer-dropup"; @@ -2189,6 +2190,9 @@ export function MessageInput({ published catalogue. */} {agentType === "atlas-agent" && } {agentType === "atlas-agent" && } + {/* Every agent's semantic_search is keyword-only until the code + model is downloaded; this is where the user can find out. */} + {projectPath && } {/* Right side, in this order: the session's usage, the agent's own knobs, then the live diff --git a/src/features/settings/lib/code-model-nudge.ts b/src/features/settings/lib/code-model-nudge.ts new file mode 100644 index 000000000..1255d5e85 --- /dev/null +++ b/src/features/settings/lib/code-model-nudge.ts @@ -0,0 +1,51 @@ +// When to suggest downloading the code search model. Without it, +// `semantic_search` quietly runs keyword + symbol only and says so only in the +// agent's tool output, which the user never reads, so plain-English questions +// miss answers with nothing on screen to explain why. + +import type { ModelStatus } from "./models-api"; + +const DISMISSED_KEY = "atlas:code-model-nudge:dismissed:v1"; + +/** The selected code model, when it is missing and the user hasn't waved the + * suggestion off for it. `null` means: show nothing. Dismissal is per model, + * so picking a different code model asks again. */ +export function codeModelToSuggest( + list: readonly ModelStatus[], + dismissed: readonly string[], +): ModelStatus | null { + const model = list.find((m) => m.kind === "code_embedding" && m.selected); + if (!model || model.downloaded || !model.compatible) return null; + return dismissed.includes(model.id) ? null : model; +} + +export function readDismissed(): string[] { + try { + const raw = localStorage.getItem(DISMISSED_KEY); + const parsed: unknown = raw ? JSON.parse(raw) : []; + return Array.isArray(parsed) ? parsed.filter((v): v is string => typeof v === "string") : []; + } catch { + return []; + } +} + +export function dismiss(id: string): string[] { + const next = [...new Set([...readDismissed(), id])]; + try { + localStorage.setItem(DISMISSED_KEY, JSON.stringify(next)); + } catch { + // Storage full or blocked: the pill still hides for this session. + } + return next; +} + +/** Overall download progress, 0–100, across a model's files. */ +export function downloadPercent(p: { + fileIndex: number; + fileCount: number; + received: number; + total: number; +}): number { + const file = p.total ? p.received / p.total : 0; + return Math.min(100, Math.round(((p.fileIndex + file) / Math.max(1, p.fileCount)) * 100)); +} diff --git a/src/features/settings/stores/models-store.ts b/src/features/settings/stores/models-store.ts index 20552b66e..50ec6f3fc 100644 --- a/src/features/settings/stores/models-store.ts +++ b/src/features/settings/stores/models-store.ts @@ -13,6 +13,7 @@ import { type ModelStatus, type DownloadProgress, } from "../lib/models-api"; +import { dismiss, readDismissed } from "../lib/code-model-nudge"; import type { UnlistenFn } from "@tauri-apps/api/event"; interface ModelsState { @@ -22,6 +23,9 @@ interface ModelsState { downloading: Record; /** model id currently being removed/selected (inline busy). */ pending: string | null; + /** Code models the user waved the composer's download nudge off for. Held + * here, not per pill, so dismissing it in one composer hides it in all. */ + codeNudgeDismissed: string[]; actions: { init: () => Promise; refresh: () => Promise; @@ -29,6 +33,8 @@ interface ModelsState { remove: (id: string) => Promise; /** Returns whether a memory re-index is required (embedding switch). */ select: (id: string) => Promise; + /** Never suggest downloading code model `id` again (persisted). */ + dismissCodeNudge: (id: string) => void; }; } @@ -41,6 +47,7 @@ export const useModelsStore = createSelectors( loaded: false, downloading: {}, pending: null, + codeNudgeDismissed: readDismissed(), actions: { init: async () => { await get().actions.refresh(); @@ -120,6 +127,8 @@ export const useModelsStore = createSelectors( set({ pending: null }); } }, + + dismissCodeNudge: (id) => set({ codeNudgeDismissed: dismiss(id) }), }, })), ); From c9e7a403510847915a043c6ed0d4290ad8d6a7c5 Mon Sep 17 00:00:00 2001 From: Uzayer Masud <16221196+uzayer@users.noreply.github.com> Date: Sat, 10 Oct 2026 00:00:46 +0600 Subject: [PATCH 3/3] fix(memory): capture recovers, the briefing tells the truth, search is bounded, and no session token reaches an agent's argv Capture and briefing - extraction catches up on recorded sessions in prompt-sized passes, at most 3 per job, skipping anything active in the last 15 minutes (in any process); failures log at warn instead of debug - sessions imported from outside Atlas (terminal agents) are read only when the new per-project 'Learn from sessions outside Atlas' setting is on; it is off by default because it sends their transcripts to the extraction model - a plan whose items are all done (or superseded) is no longer the active plan - filesChanged keeps only paths under the project root or its worktrees, with .., ~ and linked-worktree roots resolved; out-of-root paths are not captured - the briefing always carries handoff (null when there is none) and builds one for a session that never recorded its end memory_search - documents are excerpts around the matching passage inside a max_output_tokens budget (1000-12000, default 4096) with a truncation marker - a relevance floor: a match must be supported by its words or its meaning, and a query nothing supports returns an empty result with a note MCP credentials (ADR-0020, amends 0010, 0015, 0019) - only the in-process agent gets the loopback servers over HTTP with a bearer header; every ACP agent gets them through the stdio bridge, which reads the token from a 0600 file, because adapters copy an entry's headers and env onto the agent's command line - token files are removed when their session ends; ended launches are swept at startup off the runtime thread; a live launch holds a lock and is never swept Tests for each; AtlasMemBench passes unchanged; parity_bench fixtures now carry a body as the corpus format requires. Co-authored-by: Ahammad Nafiz --- CONTEXT.md | 2 +- .../atlas-agent-manager/tests/session_mcp.rs | 56 ++ crates/atlas-agent-servers/src/connection.rs | 3 + crates/atlas-agent-servers/src/session_mcp.rs | 32 +- crates/atlas-memory/src/handoff.rs | 181 +++- crates/atlas-memory/src/lib.rs | 8 +- crates/atlas-memory/src/parity_bench.rs | 9 +- crates/atlas-memory/src/record.rs | 295 ++++++- crates/atlas-memory/src/retrieve.rs | 682 ++++++++++++++- .../src/engine/connection.rs | 3 + ...emory-is-pulled-through-the-tool-server.md | 2 +- ...is-served-by-the-atlas-code-tool-server.md | 2 +- ...rries-memory-to-agents-without-http-mcp.md | 2 +- ...-an-mcp-entry-handed-to-another-process.md | 35 + docs/architecture.md | 6 +- docs/reference/configuration.md | 13 + src-tauri/src/commands/agents.rs | 23 +- src-tauri/src/commands/code_server/offers.rs | 15 +- src-tauri/src/commands/code_server/tests.rs | 60 +- src-tauri/src/commands/memory_bridge.rs | 368 +++++++- src-tauri/src/commands/memory_capture.rs | 90 ++ src-tauri/src/commands/memory_delta.rs | 122 ++- src-tauri/src/commands/memory_extract.rs | 789 +++++++++++++++++- src-tauri/src/commands/memory_indexer.rs | 139 ++- .../src/commands/memory_server/briefing.rs | 225 ++++- src-tauri/src/commands/memory_server/mod.rs | 13 +- .../src/commands/memory_server/offers.rs | 158 +++- src-tauri/src/commands/memory_server/tests.rs | 357 +++++++- src-tauri/src/commands/memory_server/tools.rs | 145 +++- src-tauri/src/commands/memory_sharing.rs | 189 ++++- src-tauri/src/commands/org_server/tests.rs | 19 +- src-tauri/src/commands/ui_server/tests.rs | 20 +- src-tauri/src/lib.rs | 30 +- src-tauri/src/main.rs | 5 +- src/dev/mock-backend/fixtures/memory.ts | 11 + .../mock-backend/scenarios/northwind-misc.ts | 6 + .../memory-sharing-controls.test.tsx | 82 ++ .../components/memory-sharing-controls.tsx | 45 +- src/features/memory/lib/memory-sharing-api.ts | 8 +- .../memory/lib/memory-sharing-copy.test.ts | 15 +- .../memory/lib/memory-sharing-copy.ts | 10 + .../memory/stores/memory-sharing-store.ts | 22 +- .../settings/components/settings-controls.tsx | 9 + 43 files changed, 3999 insertions(+), 307 deletions(-) create mode 100644 docs/adr/0020-no-credential-rides-an-mcp-entry-handed-to-another-process.md create mode 100644 src/features/memory/components/memory-sharing-controls.test.tsx diff --git a/CONTEXT.md b/CONTEXT.md index f00bd6a08..b9b391c77 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -58,7 +58,7 @@ Shared memory also records **session lifecycle** (session start and end, todo ad Terms that are *not* shared memory: **capture** (raw transcripts, see Timeline), **knowledge notes** (user-written pages in the Knowledge panel), the **codebase index** (derived from source), and any agent's own memory files (`CLAUDE.md`, `AGENTS.md`, Claude's auto-memory directory). These may be *sources* shared memory cites or imports, but an entry in them is not an entry in shared memory. -Shared memory reaches an agent one way only: it **pulls** it through the **memory tool server**, the in-process MCP server Atlas hands every session that can take it (ADR-0010). Nothing is prepended to a user's message; the text goes to the agent exactly as typed. The server's **instructions** are the protocol — read first, write as you learn — and its tools are, in that order: **`memory_briefing`** (the first look of a session: preferences, working memory, the **durable index**, the **curated pack**, the **recent-session handoff** and the **handoff note**), **`memory_changes`** (what *other* sessions recorded since this session's **last look**), **`memory_search`** (the record plus the codebase index's documents), **`memory_get`**, **`memory_list`**, **`memory_history`** (every revision of one entry), **`memory_why`** (the recorded sessions behind a file or a commit), **`memory_remember`**, **`memory_feedback`** (useful, wrong or stale) and **`memory_forget`**. An agent without HTTP MCP reaches the same server through the **stdio bridge**, the Atlas binary run as `atlas mcp-bridge` (ADR-0019). +Shared memory reaches an agent one way only: it **pulls** it through the **memory tool server**, the in-process MCP server Atlas hands every session that can take it (ADR-0010). Nothing is prepended to a user's message; the text goes to the agent exactly as typed. The server's **instructions** are the protocol — read first, write as you learn — and its tools are, in that order: **`memory_briefing`** (the first look of a session: preferences, working memory, the **durable index**, the **curated pack**, the **recent-session handoff** and the **handoff note**), **`memory_changes`** (what *other* sessions recorded since this session's **last look**), **`memory_search`** (the record plus the codebase index's documents), **`memory_get`**, **`memory_list`**, **`memory_history`** (every revision of one entry), **`memory_why`** (the recorded sessions behind a file or a commit), **`memory_remember`**, **`memory_feedback`** (useful, wrong or stale) and **`memory_forget`**. Every ACP agent, whatever it advertises, reaches the same server through the **stdio bridge**, the Atlas binary run as `atlas mcp-bridge`, its token in a private file and never in the MCP entry (ADR-0019, ADR-0020); only the in-process native agent gets it over HTTP. - The **durable index** is one capped line per durable entry, each kind's best up to its display cap, ranked by recency (two-week half-life), use count and confidence; `memory_get` expands a line. It lists shared-memory entries; it is not the codebase index, and not the retrieval index `memory_search` also searches. - The **curated pack** is the high-signal part of the project's foreign memory files (Claude's memory directory, `CLAUDE.md`, `AGENTS.md`), recency-ranked and budget-bounded. The **recent-session handoff** is the tail of the most recent other session in the scope, whichever agent ran it, read from capture (or, where capture was never enabled, from Atlas's own session transcripts) — never from an agent's own files. diff --git a/crates/atlas-agent-manager/tests/session_mcp.rs b/crates/atlas-agent-manager/tests/session_mcp.rs index f651609e5..d5a8fb26c 100644 --- a/crates/atlas-agent-manager/tests/session_mcp.rs +++ b/crates/atlas-agent-manager/tests/session_mcp.rs @@ -103,6 +103,10 @@ async fn an_agent_advertising_http_mcp_gets_the_server_with_its_token_on_session !asked[0].org_access, "an ACP connection never carries organisation access (ADR-0014), whatever the agent is", ); + assert!( + !asked[0].in_process, + "an ACP agent is another process, so no credential may ride its entries (ADR-0020)", + ); assert_eq!(asked[0].agent_id.as_str(), "fake-agent"); assert_eq!( asked[0].session_id, None, @@ -154,6 +158,58 @@ async fn an_agent_without_http_mcp_gets_no_server_entry() { drop(thread); } +/// Offers the host's server the way it reaches an ACP agent (ADR-0020): a +/// stdio bridge whose entry names a token file and holds no secret. +struct BridgeOffering; + +impl SessionMcpServers for BridgeOffering { + fn offer(&self, _request: &SessionMcpRequest) -> SessionMcpOffer { + let server = acp::McpServer::Stdio( + acp::McpServerStdio::new("atlas_memory", "/Applications/Atlas.app/atlas").args(vec![ + "mcp-bridge".into(), + "http://127.0.0.1:4321/mcp".into(), + "/private/atlas/mcp-bridge/l/t.token".into(), + ]), + ); + SessionMcpOffer::new(vec![server], |_| {}) + } +} + +/// Stdio is the one transport every ACP agent takes, advertised or not; the +/// bridge reaches the agent as offered, pre-approved for the Claude Code +/// adapter like an HTTP entry. +#[tokio::test] +async fn a_bridged_server_reaches_an_agent_without_http_mcp_pre_approved() { + let Some((manager, pid_file)) = + manager_offering_mcp("mcp-new-bridge", json!({}), Arc::new(BridgeOffering)) + else { + eprintln!("skipping: no python3 on this machine"); + return; + }; + let thread = manager + .new_session(custom("fake-agent"), vec![std::env::temp_dir()]) + .await + .expect("a session opens on the real agent"); + + let requests = session_requests(&pid_file); + assert_eq!( + requests[0].1["mcpServers"], + json!([{ + "name": "atlas_memory", + "command": "/Applications/Atlas.app/atlas", + "args": ["mcp-bridge", "http://127.0.0.1:4321/mcp", "/private/atlas/mcp-bridge/l/t.token"], + "env": [], + }]) + ); + assert_eq!( + requests[0].1["_meta"]["claudeCode"]["options"]["allowedTools"], + json!(["mcp__atlas_memory"]), + ); + + manager.shutdown(); + drop(thread); +} + #[tokio::test] async fn a_loaded_session_gets_the_server_and_its_token_is_bound_to_that_session() { let offering = Arc::new(TokenOffering::default()); diff --git a/crates/atlas-agent-servers/src/connection.rs b/crates/atlas-agent-servers/src/connection.rs index 0fb2f512b..aff5e95ed 100644 --- a/crates/atlas-agent-servers/src/connection.rs +++ b/crates/atlas-agent-servers/src/connection.rs @@ -483,6 +483,9 @@ impl AcpConnection { // Likewise for the organisation: a binary Atlas does not own // is never handed the user's organisation (ADR-0014). org_access: false, + // A separate program, free to put its MCP configuration on a + // command line: no credential may ride an entry (ADR-0020). + in_process: false, cwd: cwd.to_path_buf(), session_id: session_id.cloned(), }, diff --git a/crates/atlas-agent-servers/src/session_mcp.rs b/crates/atlas-agent-servers/src/session_mcp.rs index bed594faa..fe50e0052 100644 --- a/crates/atlas-agent-servers/src/session_mcp.rs +++ b/crates/atlas-agent-servers/src/session_mcp.rs @@ -14,6 +14,12 @@ //! sends a transport the agent did not advertise ([`admissible`]). ACP makes //! stdio mandatory for every agent and HTTP/SSE opt-in through //! `mcpCapabilities`, and that is the only gate: never the agent's identity. +//! +//! Nothing secret may ride an entry handed to an agent in another process +//! ([`SessionMcpRequest::in_process`]): an ACP adapter is free to put its MCP +//! configuration on a command line, where any local user's `ps` reads it (the +//! Claude Agent SDK passes `--mcp-config `, headers and `env` included). +//! ADR-0020. use std::path::PathBuf; use std::sync::Arc; @@ -42,6 +48,13 @@ pub struct SessionMcpRequest { /// binary is never handed the user's organisation — decided by the /// connection, never by which agent it is. pub org_access: bool, + /// Whether the agent runs **inside the Atlas process**, so an entry it is + /// handed never leaves Atlas's memory. An agent in another process (every + /// ACP agent) may write its MCP configuration anywhere — a command line, + /// a config file — so the host must not put a credential in any of an + /// entry's fields: not a header, not an argument, not `env` (ADR-0020). + /// A property of the connection, never of which agent it is. + pub in_process: bool, /// The directory the session runs in. pub cwd: PathBuf, /// The session being loaded or resumed; `None` for a new session, whose @@ -360,7 +373,8 @@ pub fn offer_for( provider.map_or_else(SessionMcpOffer::none, |p| p.offer(request)) } -/// `_meta` for a session request carrying `servers`: the Claude Code adapter +/// `_meta` for a session request carrying `servers` (only ever the host's +/// offer, never the user's own servers): the Claude Code adapter /// passes `_meta.claudeCode.options` to the Agent SDK, whose `allowedTools` rule /// `mcp__` runs that server's tools without a permission prompt: the /// standing the native agent's servers have (`default_tools_approval_mode = @@ -372,8 +386,11 @@ pub fn preapproval_meta(servers: &[acp::McpServer], ask_first: &AskFirst) -> Opt let rules: Vec = servers .iter() .filter_map(|server| match server { - acp::McpServer::Http(http) if ask_first.tools_on(&http.name).next().is_none() => { - Some(format!("mcp__{}", http.name).into()) + acp::McpServer::Http(acp::McpServerHttp { name, .. }) + | acp::McpServer::Stdio(acp::McpServerStdio { name, .. }) + if ask_first.tools_on(name).next().is_none() => + { + Some(format!("mcp__{name}").into()) } _ => None, }) @@ -414,11 +431,12 @@ mod tests { } /// Claude Code asked before every Atlas tool call (review of #354): the - /// offered HTTP servers ride the session request as pre-approved. + /// offered servers ride the session request as pre-approved, whichever + /// transport carries them (an ACP agent gets the stdio bridge, ADR-0020). #[test] - fn offered_http_servers_are_preapproved_for_claude_code() { - let stdio = acp::McpServer::Stdio(acp::McpServerStdio::new("local", "/bin/true")); - let servers = [http("atlas_memory"), http("atlas_code"), stdio]; + fn offered_servers_are_preapproved_for_claude_code() { + let stdio = acp::McpServer::Stdio(acp::McpServerStdio::new("atlas_code", "/bin/atlas")); + let servers = [http("atlas_memory"), stdio]; let meta = preapproval_meta(&servers, &AskFirst::none()).unwrap(); assert_eq!( serde_json::Value::Object(meta), diff --git a/crates/atlas-memory/src/handoff.rs b/crates/atlas-memory/src/handoff.rs index 37abb7886..860735a8a 100644 --- a/crates/atlas-memory/src/handoff.rs +++ b/crates/atlas-memory/src/handoff.rs @@ -208,10 +208,36 @@ pub fn build_handoff( _ => {} } } + // Only this scope's files: an agent here may also have edited another + // checkout, `/tmp` or `~/.claude`. + let dirs = store.scope_dirs(); + note.files.retain(|f| crate::record::path_in_dirs(f, &dirs)); note.files.truncate(MAX_FILES); Ok(note) } +/// Whether every item of `plan` is done: it has at least one +/// `[status] item` line and all of them are `completed` (or `done`, +/// `cancelled`, `skipped`, `superseded`). A finished plan is no longer the +/// active plan. +pub fn plan_is_finished(plan: &str) -> bool { + let mut items = 0usize; + for line in plan.lines() { + let line = line.trim().trim_start_matches("- "); + let Some((status, _)) = line.strip_prefix('[').and_then(|l| l.split_once(']')) else { + continue; + }; + items += 1; + if !matches!( + status.trim(), + "completed" | "complete" | "done" | "cancelled" | "canceled" | "skipped" | "superseded" + ) { + return false; + } + } + items > 0 +} + /// The plan lines still to do: `[pending]` and `[in_progress]`. fn open_items(plan: &str) -> Vec { plan.lines() @@ -232,7 +258,7 @@ fn push_unique(list: &mut Vec, item: String) { #[cfg(test)] mod tests { use super::*; - use crate::record::NewEvent; + use crate::record::{EventKind, NewEvent}; #[test] fn the_handoff_note_lists_what_the_session_left() { @@ -298,15 +324,162 @@ mod tests { assert_eq!(note.started_at, Some(1)); store.record_episode(¬e).unwrap(); assert_eq!(store.last_episode("s-b").unwrap(), Some(note.clone())); - assert_eq!( - store.last_episode("s-a").unwrap(), - None, + assert_ne!( + store.last_episode("s-a").unwrap().map(|n| n.session), + Some("s-a".to_string()), "a session is not handed its own note" ); assert_eq!(store.episodes_since(0, 10).unwrap(), vec![note]); let _ = std::fs::remove_dir_all(&root); } + fn scope(label: &str) -> (std::path::PathBuf, std::sync::Arc) { + let root = + std::env::temp_dir().join(format!("atlas-handoff-{label}-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&root).unwrap(); + let store = crate::record::open_scope(&root).unwrap(); + (root, store) + } + + fn log( + store: &RecordStore, + session: &str, + kind: EventKind, + key: &str, + payload: serde_json::Value, + ts: i64, + ) { + store + .append_event( + NewEvent { + agent: "claude-acp".into(), + session_id: session.into(), + kind, + key: key.into(), + payload, + }, + ts, + ) + .unwrap(); + } + + /// A session that edited another checkout, `/tmp` and `~/.claude` hands + /// on only the files of this scope. + #[test] + fn the_note_lists_only_files_under_the_scope_root() { + let (root, store) = scope("files"); + let root_canon = root.canonicalize().unwrap(); + let inside = root_canon.join("src/lib.rs").to_string_lossy().into_owned(); + for (i, path) in [ + inside.as_str(), + "src/relative.rs", + "/Users/someone/Developer/other-checkout/src/main.rs", + "/tmp/homebrew-cask-pr/Casks/a/atlas-ai.rb", + "~/.claude/plans/plan.md", + "../escape.rs", + ] + .into_iter() + .enumerate() + { + log( + &store, + "s-a", + EventKind::FileChanged, + path, + serde_json::json!({"path": path, "summary": "Edit"}), + i as i64 + 1, + ); + } + let note = build_handoff(&store, "s-a", "claude-acp", 10).unwrap(); + assert_eq!(note.files, ["src/relative.rs", inside.as_str()]); + let _ = std::fs::remove_dir_all(&root); + } + + /// Notes are stored only when a session's end is recorded. A session + /// that never ended (still open, crashed, ran before notes existed) is + /// still the previous session: its note is built from what it logged. + #[test] + fn a_session_that_never_ended_still_hands_off() { + let (root, store) = scope("unended"); + log( + &store, + "s-old", + EventKind::SessionStart, + "", + serde_json::json!({}), + 1, + ); + store + .remember( + crate::record::NewEntry { + kind: crate::record::EntryKind::Decision, + key: String::new(), + content: "Keep SQLite in WAL mode".into(), + source: "claude-acp".into(), + agent: "claude-acp".into(), + session_id: "s-old".into(), + confidence: 1.0, + at: 2, + }, + 2, + ) + .unwrap(); + log( + &store, + "s-old", + EventKind::FileChanged, + "src/db.rs", + serde_json::json!({"path": "src/db.rs", "summary": "Edit"}), + 3, + ); + // A session with only its start logged nothing worth handing on. + log( + &store, + "s-idle", + EventKind::SessionStart, + "", + serde_json::json!({}), + 4, + ); + + let note = store.last_episode("s-new").unwrap().expect("a handoff"); + assert_eq!(note.session, "s-old"); + assert_eq!(note.decisions, ["Keep SQLite in WAL mode"]); + assert_eq!(note.files, ["src/db.rs"]); + assert_eq!(note.ended_at, 3, "its last event"); + + // A stored note that ended later wins over an older unended session. + let stored = HandoffNote { + session: "s-ended".into(), + agent: "codex-acp".into(), + ended_at: 50, + files: vec!["src/later.rs".into(), "/tmp/elsewhere.rs".into()], + ..Default::default() + }; + store.record_episode(&stored).unwrap(); + let handed = store.last_episode("s-new").unwrap().unwrap(); + assert_eq!(handed.session, "s-ended"); + assert_eq!( + handed.files, + ["src/later.rs"], + "stored notes are filtered on read" + ); + let _ = std::fs::remove_dir_all(&root); + } + + #[test] + fn a_plan_is_finished_when_every_item_is_done() { + assert!(plan_is_finished( + "- [completed] Create branch\n- [completed] Commit, push, open PR" + )); + assert!(plan_is_finished("- [completed] A\n- [cancelled] B")); + assert!(plan_is_finished("- [completed] A\n- [superseded] B")); + assert!(!plan_is_finished("- [completed] A\n- [in_progress] B")); + assert!(!plan_is_finished("- [pending] A")); + assert!(!plan_is_finished("free text with no items")); + assert!(!plan_is_finished("")); + } + #[test] fn facts_lead_the_files_and_fill_the_recorded_fields() { let mut note = HandoffNote { diff --git a/crates/atlas-memory/src/lib.rs b/crates/atlas-memory/src/lib.rs index 4dd5d74ec..9ed855358 100644 --- a/crates/atlas-memory/src/lib.rs +++ b/crates/atlas-memory/src/lib.rs @@ -23,11 +23,11 @@ pub mod provider; pub mod store; // Step 6 (implemented): docstore (id→display-text side-map) + retrieve (HNSW, -// with the global blend). `retrieve` only adds an `impl MemoryEngine`, so it is a -// plain child module (private) — it reaches the engine's private fields as a -// descendant. +// with the global blend). `retrieve` adds an `impl MemoryEngine` — a child +// module, so it reaches the engine's private fields as a descendant — and the +// query-term relevance floor and excerpt window `memory_search` shares. pub mod docstore; -mod retrieve; +pub mod retrieve; // The corpus index behind `MemoryEngine`: documents, BM25 and cached vectors // in `corpus.sqlite`, a per-model vector file that heals from the cache. diff --git a/crates/atlas-memory/src/parity_bench.rs b/crates/atlas-memory/src/parity_bench.rs index 919d96b95..fbb2ffc6a 100644 --- a/crates/atlas-memory/src/parity_bench.rs +++ b/crates/atlas-memory/src/parity_bench.rs @@ -149,23 +149,26 @@ fn retrieve_is_agent_agnostic_and_well_formed_when_model_available() { let root = tmp.path().to_path_buf(); let mut engine = MemoryEngine::open(root); + // Indexed text is `title\n\nbody` (`docstore::split_embedded`); a doc + // with no body is returned with empty `text`, by design. let corpus = vec![ CorpusDoc { id: "d1".into(), - text: "Auth: the project uses Better Auth with database-backed sessions.".into(), + text: "Auth\n\nThe project uses Better Auth with database-backed sessions.".into(), content_hash: "h1".into(), corpus: "claude".into(), }, CorpusDoc { id: "d2".into(), - text: "Indexing: a background MemoryIndexer rebuilds the HNSW off the chat turn." + text: "Indexing\n\nA background MemoryIndexer rebuilds the HNSW off the chat turn." .into(), content_hash: "h2".into(), corpus: "codebase".into(), }, CorpusDoc { id: "d3".into(), - text: "Retrieval searches usearch HNSW and blends global memory via RRF.".into(), + text: "Retrieval\n\nRetrieval searches usearch HNSW and blends global memory via RRF." + .into(), content_hash: "h3".into(), corpus: "codebase".into(), }, diff --git a/crates/atlas-memory/src/record.rs b/crates/atlas-memory/src/record.rs index 3949e1ed5..92a5afe6b 100644 --- a/crates/atlas-memory/src/record.rs +++ b/crates/atlas-memory/src/record.rs @@ -757,6 +757,44 @@ impl RecordStore { &self.root } + /// The directories whose files are this scope's: the root and every + /// worktree of its repository — the main one and each linked one (read + /// from `/worktrees/*/gitdir`, no git process), found the + /// same way when the root is itself a linked worktree. Each in its given + /// and its canonical spelling. + pub fn scope_dirs(&self) -> Vec { + let mut dirs = vec![self.root.clone()]; + if let Some(common) = git_common_dir(&self.root) { + // A non-bare repository's common dir is its main worktree's `.git`. + if common.file_name().is_some_and(|n| n == ".git") { + if let Some(main) = common.parent() { + dirs.push(main.to_path_buf()); + } + } + if let Ok(entries) = std::fs::read_dir(common.join("worktrees")) { + for entry in entries.flatten() { + let Ok(gitdir) = std::fs::read_to_string(entry.path().join("gitdir")) else { + continue; + }; + // `gitdir` names the worktree's `.git` file. + if let Some(worktree) = Path::new(gitdir.trim()).parent() { + dirs.push(normalize_lexically(worktree)); + } + } + } + } + let canonical: Vec = dirs.iter().filter_map(|d| d.canonicalize().ok()).collect(); + dirs.extend(canonical); + dirs.sort(); + dirs.dedup(); + dirs + } + + /// Whether `path` names a file of this scope (see [`path_in_dirs`]). + pub fn in_scope(&self, path: &str) -> bool { + path_in_dirs(path, &self.scope_dirs()) + } + pub(crate) fn conn(&self) -> MutexGuard<'_, Connection> { self.conn .lock() @@ -857,6 +895,22 @@ impl RecordStore { .optional()?) } + /// Each session's newest event at or after `since`, by session id: when + /// it last did anything this scope's log saw. Sessions with no event + /// since then are absent. + pub fn last_event_by_session( + &self, + since: i64, + ) -> Result> { + let conn = self.conn(); + let mut stmt = conn.prepare( + "SELECT session, MAX(ts) FROM events WHERE session <> '' AND ts >= ?1 \ + GROUP BY session", + )?; + let rows = stmt.query_map([since], |r| Ok((r.get(0)?, r.get(1)?)))?; + Ok(rows.collect::>()?) + } + /// Up to `limit` events, newest first. pub fn events_newest(&self, limit: usize) -> Result> { let conn = self.conn(); @@ -895,24 +949,88 @@ impl RecordStore { Ok(()) } - /// The newest handoff note left by any session but `except_session`, - /// holding only what memory still trusts (see `live_note`). + /// The newest handoff note of any session but `except_session`, holding + /// only what memory still trusts (see `live_note`) and only this scope's + /// files. + /// + /// A note is stored when a session's end is recorded, and many sessions + /// never get one: one still open in another tab, one cut off by a crash, + /// one that ran before notes existed. So when a session without a note + /// logged work after the newest stored note ended, its note is built now + /// from its events ([`crate::handoff::build_handoff`]), `ended_at` being + /// its last event, and that one is handed on instead. pub fn last_episode( &self, except_session: &str, ) -> Result> { - let conn = self.conn(); - let note: Option = conn + let stored: Option<(String, i64)> = self + .conn() .query_row( - "SELECT note FROM episodes WHERE session <> ?1 ORDER BY ended_at DESC LIMIT 1", + "SELECT note, ended_at FROM episodes WHERE session <> ?1 \ + ORDER BY ended_at DESC LIMIT 1", [except_session], - |r| r.get(0), + |r| Ok((r.get(0)?, r.get(1)?)), ) .optional()?; - match note { - Some(n) => live_note(&conn, &n), - None => Ok(None), + let since = stored.as_ref().map_or(i64::MIN, |(_, at)| *at); + let mut note = match self.unrecorded_handoff(except_session, since)? { + Some(built) => Some(built), + None => match stored { + Some((raw, _)) => live_note(&self.conn(), &raw)?, + None => None, + }, + }; + if let Some(note) = &mut note { + let dirs = self.scope_dirs(); + note.files + .retain(|f| path_in_dirs(f.strip_suffix(" (deleted)").unwrap_or(f), &dirs)); + } + Ok(note) + } + + /// The note of the newest session other than `except_session` that has + /// no stored note and logged work (anything but its start and end) after + /// `since`, built from its events; `None` when there is none, or when the + /// few newest such sessions left nothing to hand on. + fn unrecorded_handoff( + &self, + except_session: &str, + since: i64, + ) -> Result> { + /// How many such sessions are tried, newest first. + const TRIED: i64 = 5; + let candidates: Vec<(String, String, i64, Option)> = { + let conn = self.conn(); + let mut stmt = conn.prepare(&format!( + "SELECT e.session, MAX(e.agent), MAX(e.ts) AS last, \ + (SELECT s.ended_at FROM sessions s WHERE s.session_id = e.session) \ + FROM events e {LIVE_EVENTS} \ + AND e.session <> '' AND e.session <> ?1 \ + AND e.kind NOT IN ('session_start', 'session_end') \ + AND e.session NOT IN (SELECT session FROM episodes) \ + GROUP BY e.session HAVING last > ?2 \ + ORDER BY last DESC LIMIT ?3" + ))?; + let rows = stmt.query_map(params![except_session, since, TRIED], |r| { + Ok((r.get(0)?, r.get(1)?, r.get(2)?, r.get(3)?)) + })?; + rows.collect::>>()? + }; + for (session, agent, last, ended_at) in candidates { + let built = crate::handoff::build_handoff( + self, + &session, + &agent, + ended_at.unwrap_or(last).max(last), + )?; + let Some(note) = live_note(&self.conn(), &serde_json::to_string(&built)?)? else { + continue; + }; + if !note.is_empty() { + return Ok(Some(note)); + } } + Ok(None) } /// Handoff notes that ended after `since`, oldest first, at most `limit`, @@ -1861,6 +1979,101 @@ pub struct Source { pub commits: Vec, } +/// Whether `path` lies in one of `dirs`: a relative path that does not climb +/// out (`..`) always does; an absolute one (or `~/…`, read from `$HOME`) +/// only under one of them, compared component by component after `.` and +/// `..` are resolved, as given and canonicalised. An agent's edits to +/// another checkout, `/tmp` or `~/.claude` are not this scope's files, and +/// `/../elsewhere` is not under the root. +pub fn path_in_dirs(path: &str, dirs: &[PathBuf]) -> bool { + let path = path.trim(); + if path.is_empty() { + return false; + } + let expanded; + let p = if let Some(rest) = path.strip_prefix('~') { + // `~user/…` names someone else's home: never this scope's. + let rest = match rest.strip_prefix(['/', '\\']) { + Some(rest) => rest, + None if rest.is_empty() => "", + None => return false, + }; + let Some(home) = home_dir() else { + return false; + }; + expanded = home.join(rest); + expanded.as_path() + } else { + Path::new(path) + }; + let p = normalize_lexically(p); + if !p.has_root() { + return !p + .components() + .any(|c| matches!(c, std::path::Component::ParentDir)); + } + let dirs: Vec = dirs.iter().map(|d| normalize_lexically(d)).collect(); + if dirs.iter().any(|d| p.starts_with(d)) { + return true; + } + // A symlinked spelling (`/tmp` → `/private/tmp`): the file, else its + // directory, canonicalised. A deleted file's directory usually remains. + let canonical = p.canonicalize().ok().or_else(|| { + let parent = p.parent()?.canonicalize().ok()?; + Some(parent.join(p.file_name()?)) + }); + canonical.is_some_and(|c| dirs.iter().any(|d| c.starts_with(d))) +} + +/// The user's home directory, from the environment. +fn home_dir() -> Option { + let var = if cfg!(windows) { "USERPROFILE" } else { "HOME" }; + std::env::var_os(var) + .filter(|h| !h.is_empty()) + .map(PathBuf::from) +} + +/// `path` with `.` dropped and each `..` taking off the component before it, +/// without touching the filesystem. A `..` with nothing to take off stays +/// (relative) or is dropped (at the root, where `/..` is `/`). +fn normalize_lexically(path: &Path) -> PathBuf { + use std::path::Component; + let mut out: Vec> = Vec::new(); + for c in path.components() { + match c { + Component::CurDir => {} + Component::ParentDir => match out.last() { + Some(Component::Normal(_)) => { + out.pop(); + } + Some(Component::RootDir | Component::Prefix(_)) => {} + _ => out.push(c), + }, + c => out.push(c), + } + } + out.iter().collect() +} + +/// The repository's common git directory for a checkout at `root`: its +/// `.git` directory, or — for a linked worktree, whose `.git` is a file +/// naming `/worktrees/` — the directory that file's +/// `commondir` points back to. Read from the files, no git process. +fn git_common_dir(root: &Path) -> Option { + let dot_git = root.join(".git"); + if dot_git.is_dir() { + return Some(dot_git); + } + let pointer = std::fs::read_to_string(&dot_git).ok()?; + let gitdir = pointer.trim().strip_prefix("gitdir:")?.trim(); + let gitdir = root.join(gitdir); + let common = match std::fs::read_to_string(gitdir.join("commondir")) { + Ok(common) => gitdir.join(common.trim()), + Err(_) => gitdir, + }; + Some(normalize_lexically(&common)) +} + /// A stable reference to the session a memory came from, in the spirit of /// the Agent Memory Repo `source` key. Opaque: Atlas has no URL scheme. pub fn source_uri(s: &Source) -> String { @@ -4082,6 +4295,70 @@ fn insert_entry(tx: &Transaction<'_>, e: &NewEntry) -> Result { pub(crate) mod tests { use super::*; + #[test] + fn a_path_that_climbs_out_of_the_root_is_not_in_it() { + let dirs = [PathBuf::from("/work/repo")]; + assert!(path_in_dirs("/work/repo/src/a.rs", &dirs)); + assert!(path_in_dirs("/work/repo/./src/../src/a.rs", &dirs)); + assert!(!path_in_dirs("/work/repo/../other/x.rs", &dirs)); + assert!(!path_in_dirs("/work/repo/src/../../other/x.rs", &dirs)); + assert!(path_in_dirs("src/../lib.rs", &dirs)); + assert!(!path_in_dirs("src/../../lib.rs", &dirs)); + // A root spelled with a `..` of its own still contains its files. + assert!(path_in_dirs( + "/work/repo/a.rs", + &[PathBuf::from("/work/x/../repo")] + )); + } + + #[test] + fn a_home_relative_path_is_read_from_home() { + let Some(home) = home_dir() else { + return; + }; + let inside = home.join("atlas-scope-test-repo"); + assert!(path_in_dirs( + "~/atlas-scope-test-repo/src/a.rs", + std::slice::from_ref(&inside) + )); + let dirs = std::slice::from_ref(&inside); + assert!(!path_in_dirs("~/.claude/plans/plan.md", dirs)); + assert!(!path_in_dirs("~/atlas-scope-test-repo/../x", dirs)); + assert!(!path_in_dirs( + "~someone/atlas-scope-test-repo/a.rs", + &[inside] + )); + } + + /// A linked worktree as the scope root still sees the main worktree and + /// the other linked ones. + #[test] + fn a_linked_worktree_root_finds_its_siblings() { + let base = tempfile::TempDir::new().unwrap(); + let main = base.path().join("main"); + let (wt_a, wt_b) = (base.path().join("wt-a"), base.path().join("wt-b")); + for (name, wt) in [("wt-a", &wt_a), ("wt-b", &wt_b)] { + let admin = main.join(".git").join("worktrees").join(name); + std::fs::create_dir_all(&admin).unwrap(); + std::fs::create_dir_all(wt).unwrap(); + std::fs::write( + admin.join("gitdir"), + format!("{}\n", wt.join(".git").display()), + ) + .unwrap(); + std::fs::write(admin.join("commondir"), "../..\n").unwrap(); + std::fs::write(wt.join(".git"), format!("gitdir: {}\n", admin.display())).unwrap(); + } + let store = RecordStore::open(&wt_a).unwrap(); + let dirs = store.scope_dirs(); + for dir in [&main, &wt_a, &wt_b] { + assert!(dirs.contains(dir), "{} in {dirs:?}", dir.display()); + } + assert!(store.in_scope(&wt_b.join("src/x.rs").to_string_lossy())); + assert!(store.in_scope(&main.join("src/x.rs").to_string_lossy())); + assert!(!store.in_scope(&base.path().join("other/x.rs").to_string_lossy())); + } + pub(crate) fn temp_root(label: &str) -> PathBuf { let dir = std::env::temp_dir().join(format!("atlas-record-{label}-{}", uuid::Uuid::new_v4())); diff --git a/crates/atlas-memory/src/retrieve.rs b/crates/atlas-memory/src/retrieve.rs index 942ec0675..7dd245a07 100644 --- a/crates/atlas-memory/src/retrieve.rs +++ b/crates/atlas-memory/src/retrieve.rs @@ -20,6 +20,15 @@ //! (#89): what it held (the legacy shared log) lives in the record store, whose //! entries are in the HNSW corpus. //! +//! 4. **Relevance floor.** RRF ranks; it never says "nothing here". So a hit +//! is admitted only on evidence of its own: its text carries query terms +//! worth at least [`MIN_TERM_COVERAGE`] of the query's IDF weight +//! ([`QueryTerms`]), or its meaning stands out of the corpus — a cosine at +//! least [`DENSE_OUTLIER_Z`] standard deviations above the query's mean +//! similarity to every doc. The z-score, not a raw cosine, because a raw +//! cosine is model-specific (bge puts unrelated text at 0.5–0.7, MiniLM at +//! 0.1–0.3). A query with no evidence anywhere returns nothing. +//! //! HyDE / lexical query expansion (the expensive full-Hybrid path) is left behind //! the off-by-default [`ENABLE_HYDE_EXPANSION`] flag — not implemented here. @@ -55,6 +64,34 @@ const LOCAL_SPARSE_THRESHOLD: usize = 3; /// near-duplicate of one already kept and dropped. const JACCARD_DUP_THRESHOLD: f32 = 0.8; +/// The share of a query's IDF weight a hit's text must carry to be admitted on +/// words alone. A third: an entry that shares only `repository` with "grep +/// index prefilter large repositories" carries ~0.1 of it and is dropped; one +/// that names the rare terms (`grep`, `prefilter`) clears it without needing +/// every word of a chatty question. +pub const MIN_TERM_COVERAGE: f32 = 1.0 / 3.0; +/// The share a document needs on words alone when the query was embedded and +/// the dense leg did not find it at all: one signal, so a stronger one. A +/// transcript's passage that happens to say `large` and `repository` carries +/// ~0.35 of "grep index prefilter large repositories" and nothing about it. +/// With no embedding there is nothing to corroborate with, and a third holds. +pub const MIN_UNCORROBORATED_COVERAGE: f32 = 0.5; +/// How far (in standard deviations of the query's similarity to the whole +/// corpus) a hit found only by meaning must stand out to be admitted. Measured +/// on a live bge-base store: the best dense hit of a query with no answer sat +/// at z 2.2–3.6, of one with an answer at z 4.1–6.4. +pub const DENSE_OUTLIER_Z: f32 = 4.0; +/// The passage a document's words are judged on by the relevance floor +/// (about 300 tokens, `memory_search`'s excerpt size). A query's terms +/// scattered across a long transcript are not an answer. +pub const PASSAGE_BYTES: usize = 1200; +/// Below this many embedded docs a mean and deviation say nothing, so the +/// dense leg is admitted on the cosine floor alone, as before the floor. +const MIN_BACKGROUND_DOCS: usize = 30; +/// The most corpus vectors read to estimate the background. A larger corpus +/// samples its best-scoring part, which raises the mean: stricter, never laxer. +const BACKGROUND_SAMPLE_MAX: usize = 20_000; + /// Off-by-default flag for HyDE / lexical query expansion (the full 183s/Q Hybrid /// path). Intentionally unimplemented in Step 6 — wired in a later step. Marked /// `allow(dead_code)` so the seam is visible without tripping the linter. @@ -111,21 +148,38 @@ impl MemoryEngine { let pool = limit.saturating_mul(4).max(20); // ── 1. Embedding (primary) ──────────────────────────────────────────── - let embed_ranked = match qvec.map(|v| self.vector_candidates(v, pool)) { - Some(Ok(r)) => r, - Some(Err(e)) => { - tracing::debug!(target: "atlas_memory::retrieve", "embedding recall failed: {e}"); - Vec::new() - } - None => Vec::new(), - }; + let (embed_ranked, outliers) = qvec + .map(|v| self.vector_candidates(v, pool)) + .unwrap_or_default(); // ── 1b. Keywords (BM25) ─────────────────────────────────────────────── let bm25_ranked = self.bm25_candidates(query, pool); - // ── 2. Jaccard dedup → top-`limit` ──────────────────────────────────── + // ── 2. Floor → Jaccard dedup → top-`limit` ──────────────────────────── + let terms = self.query_terms(query); + let dense: HashSet<&str> = embed_ranked.iter().map(|r| r.id.as_str()).collect(); + let embedded_query = qvec.is_some(); + let admitted = |fused: Vec<(RetrievedDoc, f32)>| -> Vec<(RetrievedDoc, f32)> { + fused + .into_iter() + .filter(|(d, _)| { + let by_meaning = dense.contains(d.id.as_str()) + && outliers.as_ref().is_none_or(|o| o.contains(&d.id)); + let corroborated = !embedded_query || dense.contains(d.id.as_str()); + let min = if corroborated { + MIN_TERM_COVERAGE + } else { + MIN_UNCORROBORATED_COVERAGE + }; + by_meaning || terms.supports_passage(&d.title, &d.text, min) + }) + .collect() + }; let local = jaccard_dedup( - rrf_fuse_weighted(&[(&embed_ranked, W_EMBED), (&bm25_ranked, W_BM25)]), + admitted(rrf_fuse_weighted(&[ + (&embed_ranked, W_EMBED), + (&bm25_ranked, W_BM25), + ])), limit, ); @@ -144,17 +198,47 @@ impl MemoryEngine { (&bm25_ranked, W_BM25), (&global_ranked, W_GLOBAL), ]); - jaccard_dedup(fused, limit) + jaccard_dedup(admitted(fused), limit) } /// Cosine hits for an embedded query that clear the floor, ranked best - /// first and resolved to display docs through the corpus index. - fn vector_candidates(&self, qvec: &[f32], pool: usize) -> anyhow::Result> { - let hits = self.corpus.search_dense(qvec, pool); - Ok(apply_cosine_floor(hits, COSINE_FLOOR) + /// first and resolved to display docs through the corpus index — plus the + /// ids among them that stand out of the corpus by [`DENSE_OUTLIER_Z`] + /// (`None` when the corpus is too small to say: then every one does). + fn vector_candidates( + &self, + qvec: &[f32], + pool: usize, + ) -> (Vec, Option>) { + let embedded = self.corpus.len(); + if embedded < MIN_BACKGROUND_DOCS { + let hits = self.corpus.search_dense(qvec, pool); + let ranked = apply_cosine_floor(hits, COSINE_FLOOR) + .into_iter() + .filter_map(|(id, _sim)| self.ranked(id)) + .collect(); + return (ranked, None); + } + let mut all = self + .corpus + .search_dense(qvec, embedded.min(BACKGROUND_SAMPLE_MAX)); + let outliers = dense_outliers(&all, DENSE_OUTLIER_Z); + all.truncate(pool); + let ranked = apply_cosine_floor(all, COSINE_FLOOR) .into_iter() .filter_map(|(id, _sim)| self.ranked(id)) - .collect()) + .collect(); + (ranked, Some(outliers)) + } + + /// The query's terms, weighted by their rarity in the corpus. + fn query_terms(&self, query: &str) -> QueryTerms { + let docs = self.corpus.len(); + QueryTerms::new(query).weighted(docs, |word| { + self.corpus + .search_bm25(word, docs) + .map_or(0, |ids| ids.len()) + }) } /// Docs whose words match the query (BM25), best first. @@ -316,6 +400,330 @@ fn split_memory_text(content: &str) -> (String, String) { (title, body) } +/// The ids of `hits` (`(id, cosine)`, a sample of the corpus) whose cosine is +/// at least `z` standard deviations above the sample's mean. +fn dense_outliers(hits: &[(String, f32)], z: f32) -> HashSet { + let n = hits.len() as f32; + if hits.is_empty() { + return HashSet::new(); + } + let mean = hits.iter().map(|(_, s)| *s).sum::() / n; + let sd = (hits.iter().map(|(_, s)| (*s - mean).powi(2)).sum::() / n).sqrt(); + if sd <= f32::EPSILON { + return HashSet::new(); + } + hits.iter() + .filter(|(_, s)| (*s - mean) / sd >= z) + .map(|(id, _)| id.clone()) + .collect() +} + +// ── Query terms: the relevance floor and the excerpt window ───────────────── + +/// Words that carry no topic, left out of a query's terms so "how does the +/// memory server time out" is judged on `memory`, `server`, `time`. +const STOPWORDS: &[&str] = &[ + "about", "again", "all", "also", "and", "any", "are", "been", "being", "both", "but", "can", + "could", "did", "does", "doing", "done", "each", "for", "from", "get", "gets", "got", "had", + "has", "have", "having", "her", "here", "his", "how", "into", "its", "just", "may", "might", + "more", "most", "must", "nor", "not", "now", "off", "once", "only", "onto", "other", "our", + "out", "over", "own", "same", "shall", "she", "should", "some", "such", "than", "that", "the", + "their", "them", "then", "there", "these", "they", "this", "those", "too", "under", "very", + "was", "were", "what", "when", "where", "which", "while", "who", "whom", "whose", "why", + "will", "with", "would", "yet", "you", "your", +]; + +/// `word`'s comparison key: lowercased, one inflection off (`-ies`→`-y`, +/// `-ing`, `-ed`, `-es`, `-s`, then a final `-e`), first six letters — so +/// `repositories`/`repository`, `handled`/`handle`, `tokens`/`token` meet, +/// and `prefilter`/`prefix` do not. +pub fn term_stem(word: &str) -> String { + let w = word.to_lowercase(); + let long = |b: &str| b.chars().count() >= 3; + let mut base: String = w.clone(); + if let Some(b) = w.strip_suffix("ies").filter(|b| long(b)) { + base = format!("{b}y"); + } else if !w.ends_with("ss") { + for suffix in ["ing", "ed", "es", "s"] { + if let Some(b) = w.strip_suffix(suffix).filter(|b| long(b)) { + base = b.to_string(); + break; + } + } + } + if let Some(b) = base.strip_suffix('e').filter(|b| long(b)) { + base = b.to_string(); + } + base.chars().take(6).collect() +} + +/// `text`'s words with their byte spans, split on anything not alphanumeric +/// (so `prompt_too_large` is three words). +fn words_with_spans(text: &str) -> impl Iterator { + let mut out = Vec::new(); + let mut start: Option = None; + for (i, c) in text.char_indices() { + match (c.is_alphanumeric(), start) { + (true, None) => start = Some(i), + (false, Some(s)) => { + out.push((s, i, &text[s..i])); + start = None; + } + _ => {} + } + } + if let Some(s) = start { + out.push((s, text.len(), &text[s..])); + } + out.into_iter() +} + +/// A query's distinct topic terms and what each is worth: the relevance +/// floor ([`supports`](Self::supports)) and the excerpt window +/// ([`excerpt`](Self::excerpt)) of `memory_search`. +/// +/// Terms are the query's words of three letters or more (all of them when it +/// has none that long, as the FTS query does), less [`STOPWORDS`], compared by +/// [`term_stem`]. Unweighted, every term counts the same; [`weighted`] +/// (Self::weighted) gives each its BM25 IDF, so a word every entry uses +/// (`repository`) counts for little and one nothing uses counts for a lot. +#[derive(Debug, Clone)] +pub struct QueryTerms { + /// The first spelling of each term in the query (what the FTS is asked). + words: Vec, + stems: Vec, + weights: Vec, +} + +impl QueryTerms { + pub fn new(query: &str) -> Self { + let all: Vec<&str> = words_with_spans(query).map(|(_, _, w)| w).collect(); + let long: Vec<&str> = all + .iter() + .copied() + .filter(|w| w.chars().count() >= 3) + .collect(); + let chosen = if long.is_empty() { all } else { long }; + let mut words = Vec::new(); + let mut stems: Vec = Vec::new(); + for w in chosen { + if STOPWORDS.contains(&w.to_lowercase().as_str()) { + continue; + } + let stem = term_stem(w); + if !stems.contains(&stem) { + stems.push(stem); + words.push(w.to_string()); + } + } + let weights = vec![1.0; stems.len()]; + Self { + words, + stems, + weights, + } + } + + /// Weigh each term by its BM25 IDF over `docs` documents, `df(word)` of + /// which contain it: `ln(1 + (N − df + ½) / (df + ½))`, always positive. + /// + /// A term no document uses is weighed as one used once, the rarest the + /// corpus can show. BM25 gives `df = 0` the largest weight of all, which + /// is harmless in a score (no document collects it) but wrong in a share: + /// against a small record, the words of a paraphrase nothing uses + /// ("database decision") would outweigh the one word that does match + /// ("Postgres") — in a one-entry record, by 5 to 1. + #[must_use] + pub fn weighted(mut self, docs: usize, df: impl Fn(&str) -> usize) -> Self { + self.weights = self + .words + .iter() + .map(|w| { + let df = (df(w) as f32).max(1.0); + let n = (docs as f32).max(df); + (1.0 + (n - df + 0.5) / (df + 0.5)).ln() + }) + .collect(); + self + } + + /// Whether the query has no topic term (then nothing can be judged). + pub fn is_empty(&self) -> bool { + self.stems.is_empty() + } + + /// The query's topic terms, as compared ([`term_stem`]). + pub fn stems(&self) -> &[String] { + &self.stems + } + + fn index_of(&self, word: &str) -> Option { + if word.is_empty() { + return None; + } + let stem = term_stem(word); + self.stems.iter().position(|s| *s == stem) + } + + /// The share (0–1) of the query's term weight that `text` carries. + pub fn coverage(&self, text: &str) -> f32 { + let total: f32 = self.weights.iter().sum(); + if total <= 0.0 { + return 0.0; + } + let mut seen = vec![false; self.stems.len()]; + for (_, _, w) in words_with_spans(text) { + if let Some(i) = self.index_of(w) { + seen[i] = true; + } + } + seen.iter() + .zip(&self.weights) + .filter(|(hit, _)| **hit) + .map(|(_, w)| w) + .sum::() + / total + } + + /// Whether `text` is evidence enough to be shown as a match: it carries + /// [`MIN_TERM_COVERAGE`] of the query's weight. A query with no topic term + /// cannot be judged, so everything passes. + pub fn supports(&self, text: &str) -> bool { + self.meets(text, MIN_TERM_COVERAGE) + } + + /// Whether `text` carries at least `min` of the query's weight. + fn meets(&self, text: &str, min: f32) -> bool { + // The epsilon: one term of three equally weighted is exactly a third, + // and an f32 division can land a hair under it. + self.is_empty() || self.coverage(text) + 1e-6 >= min + } + + /// The `room`-byte span of `text` (byte offsets) whose words carry the + /// most of the query's weight; the head when no term occurs. + fn best_window(&self, text: &str, room: usize) -> (usize, usize) { + let hits: Vec<(usize, usize, usize)> = words_with_spans(text) + .filter_map(|(s, e, w)| self.index_of(w).map(|i| (s, e, i))) + .collect(); + // The run of hits that fits in `room` and is worth the most. + let (mut from, mut to) = (0, room.min(text.len())); + let mut best = 0.0f32; + let mut counts = vec![0usize; self.stems.len()]; + let mut worth = 0.0f32; + let mut left = 0; + for right in 0..hits.len() { + let (_, end, i) = hits[right]; + if counts[i] == 0 { + worth += self.weights[i]; + } + counts[i] += 1; + while end - hits[left].0 > room { + let j = hits[left].2; + counts[j] -= 1; + if counts[j] == 0 { + worth -= self.weights[j]; + } + left += 1; + } + if worth > best { + best = worth; + let span = end - hits[left].0; + let start = hits[left].0.saturating_sub((room - span) / 2); + let stop = (start + room).min(text.len()); + (from, to) = (stop.saturating_sub(room), stop); + } + } + (from, to) + } + + /// Whether `title` plus the best [`PASSAGE_BYTES`] of `text` carry at + /// least `min` of the query's weight. A whole session transcript uses + /// nearly every word there is, so a share counted over all of it admits + /// it for any query; the terms have to meet where it would be excerpted. + pub fn supports_passage(&self, title: &str, text: &str, min: f32) -> bool { + if self.is_empty() || text.len() <= PASSAGE_BYTES { + return self.meets(&format!("{title} {text}"), min); + } + let (mut from, mut to) = self.best_window(text, PASSAGE_BYTES); + while !text.is_char_boundary(from) { + from += 1; + } + while !text.is_char_boundary(to) { + to -= 1; + } + self.meets(&format!("{title} {}", &text[from..to]), min) + } + + /// At most `max_bytes` of `text`, from the window that carries the most of + /// the query's weight (the head when no term occurs), cut at word + /// boundaries and marked `…` where text was dropped. `None` when the whole + /// text fits. + pub fn excerpt(&self, text: &str, max_bytes: usize) -> Option { + const MARK: &str = "…"; + if text.len() <= max_bytes { + return None; + } + let room = max_bytes.saturating_sub(2 * MARK.len()).max(1); + let (mut from, mut to) = self.best_window(text, room); + // Whole words only: start after a space, stop before one. + while !text.is_char_boundary(from) { + from += 1; + } + while !text.is_char_boundary(to) { + to -= 1; + } + if from > 0 { + if let Some(sp) = text[from..to].find(char::is_whitespace) { + if sp < 64 { + from += sp; + } + } + } + if to < text.len() { + if let Some(sp) = text[from..to].rfind(char::is_whitespace) { + if to - (from + sp) < 64 { + to = from + sp; + } + } + } + let body = text[from..to].trim(); + Some(format!( + "{}{body}{}", + if from > 0 { MARK } else { "" }, + if to < text.len() { MARK } else { "" } + )) + } +} + +/// [`QueryTerms`] for searching the shared-memory record, weighted by how +/// many of its entries use each term. +pub fn record_query_terms(store: &crate::record::RecordStore, query: &str) -> QueryTerms { + let terms = QueryTerms::new(query); + if terms.is_empty() { + return terms; + } + let conn = store.conn(); + let count = |sql: &str, arg: Option<&str>| -> usize { + let n: rusqlite::Result = match arg { + Some(a) => conn.query_row(sql, [a], |r| r.get(0)), + None => conn.query_row(sql, [], |r| r.get(0)), + }; + n.map_or(0, |n| n.max(0) as usize) + }; + let docs = count( + "SELECT COUNT(*) FROM entries WHERE state != 'archived'", + None, + ); + terms.weighted(docs, |word| { + crate::record::fts_query(word).map_or(0, |m| { + count( + "SELECT COUNT(*) FROM entries_fts WHERE entries_fts MATCH ?1", + Some(&m), + ) + }) + }) +} + /// Stable (process-independent enough) hash of a string for synthetic global ids. fn stable_hash(s: &str) -> u64 { let mut h = DefaultHasher::new(); @@ -343,6 +751,97 @@ mod tests { } } + #[test] + fn term_stems_meet_across_inflections() { + for (a, b) in [ + ("repositories", "repository"), + ("handled", "handle"), + ("tokens", "token"), + ("minutes", "minute"), + ("configured", "configure"), + ("signing", "signed"), + ] { + assert_eq!(term_stem(a), term_stem(b), "{a} / {b}"); + } + assert_ne!(term_stem("prefilter"), term_stem("prefix")); + assert_ne!(term_stem("class"), term_stem("cla")); + } + + /// The live miss: "grep index prefilter large repositories" returned a + /// branch fact whose only shared word was `repository`, used everywhere. + #[test] + fn coverage_is_weighted_by_rarity() { + let terms = + QueryTerms::new("grep index prefilter large repositories").weighted(200, |w| match w { + "grep" => 3, + "index" => 20, + "prefilter" => 0, + "large" => 10, + _ => 120, + }); + assert!(!terms.supports("The repository has a local branch named '0.3.4'")); + assert!(!terms.supports("prompt_too_large vs request_too_large error")); + assert!(terms.supports("The grep prefilter skips files that cannot match")); + // Stopwords are not terms; a query of only stopwords judges nothing. + assert_eq!(QueryTerms::new("how does the").stems().len(), 0); + assert!(QueryTerms::new("how does the").supports("anything")); + } + + /// AtlasMemBench's clobber probe: a one-entry record, and a query whose + /// other words nothing uses. The word it shares still counts. + #[test] + fn unused_query_words_do_not_drown_the_one_that_matches() { + let terms = QueryTerms::new("database decision Postgres") + .weighted(1, |w| usize::from(w == "Postgres")); + assert!(terms.supports("db Use Postgres")); + // One term in five, as in the live miss, is still not enough. + let five = + QueryTerms::new("grep index prefilter large repositories").weighted(4, |w| match w { + "large" => 1, + "repositories" => 3, + _ => 0, + }); + assert!(!five.supports("prompt_too_large 413 is a recoverable signal")); + } + + /// A long transcript that uses the query's words far apart is not an + /// answer; one where they meet in a passage is. + #[test] + fn a_long_document_is_judged_on_one_passage() { + let filler = "lorem ipsum dolor sit amet ".repeat(200); + let terms = QueryTerms::new("grep prefilter large repositories"); + let scattered = format!("grep {filler} prefilter {filler} large {filler}"); + assert!( + terms.supports(&scattered), + "the whole text carries them all" + ); + assert!(!terms.supports_passage("session", &scattered, MIN_TERM_COVERAGE)); + let together = format!("{filler} the grep prefilter for large repos {filler}"); + assert!(terms.supports_passage("session", &together, MIN_TERM_COVERAGE)); + } + + /// The excerpt is the window around the matching passage, not the head, + /// and stays inside its budget. + #[test] + fn the_excerpt_is_the_matching_passage() { + let filler = "lorem ipsum dolor sit amet ".repeat(400); + let text = format!("HEAD {filler}the grep prefilter skips files in large repos {filler}"); + let terms = QueryTerms::new("grep prefilter large repositories"); + let cut = terms.excerpt(&text, 600).expect("too long to keep whole"); + assert!(cut.len() <= 600, "{} bytes", cut.len()); + assert!( + cut.contains("grep prefilter skips files in large repos"), + "{cut}" + ); + assert!(cut.starts_with('…') && cut.ends_with('…'), "{cut}"); + assert!(!cut.contains("HEAD")); + // No term at all: the head. + let none = QueryTerms::new("sourdough").excerpt(&text, 600).unwrap(); + assert!(none.starts_with("HEAD"), "{none}"); + // Short enough: kept whole. + assert_eq!(terms.excerpt("grep prefilter", 600), None); + } + /// The cosine floor drops sub-0.30 hits BEFORE fusion ever sees them. #[test] fn cosine_floor_drops_below_threshold() { @@ -664,3 +1163,154 @@ mod fixture_corpus { std::fs::remove_dir_all(&global).ok(); } } + +/// The relevance floor over a corpus big enough to have a background (40 +/// docs): every doc sits at cosine 0.5 from every other, so a query can be +/// near all of them at once (no answer) or stand out near one (an answer). +#[cfg(test)] +mod relevance_floor { + use crate::{MemoryEngine, DIM}; + use std::path::PathBuf; + + fn tmp(name: &str) -> PathBuf { + let p = + std::env::temp_dir().join(format!("atlas-memory-floor-{}-{name}", std::process::id())); + let _ = std::fs::remove_dir_all(&p); + std::fs::create_dir_all(&p).unwrap(); + p + } + + fn vec_of(parts: &[(usize, f32)]) -> Vec { + let mut v = vec![0.0f32; DIM]; + for (axis, w) in parts { + v[*axis] += *w; + } + let n = v.iter().map(|x| x * x).sum::().sqrt(); + v.iter().map(|x| x / n).collect() + } + + /// Doc `i` (1–40) is about `topic`, mentions `repository` like every + /// other, and embeds on the shared axis 0 plus its own axis `i`. Doc 7 + /// also names the grep `prefilter`. + fn engine(root: &std::path::Path) -> MemoryEngine { + let mut engine = MemoryEngine::open(root.to_path_buf()); + let docs: Vec<(crate::CorpusDoc, Vec)> = (1..=40) + .map(|i| { + let extra = if i == 7 { + " the grep prefilter skips files" + } else { + "" + }; + ( + crate::CorpusDoc { + id: format!("doc:{i}"), + text: format!("Note {i}\n\ntopic{i} in this repository{extra}"), + content_hash: format!("h{i}"), + corpus: "test".into(), + }, + vec_of(&[(0, 1.0), (i, 1.0)]), + ) + }) + .collect(); + engine.add_embedded(&docs).unwrap(); + engine + } + + fn ids( + engine: &MemoryEngine, + query: &str, + qvec: &[f32], + global: &std::path::Path, + ) -> Vec { + engine + .retrieve_with_vector(query, Some(qvec), 5, global) + .into_iter() + .map(|d| d.id) + .collect() + } + + /// A query near every doc alike and sharing none of their words has no + /// answer: nothing is returned, though every doc clears the cosine floor. + #[test] + fn a_query_with_no_answer_returns_nothing() { + let (root, global) = (tmp("none"), tmp("none-global")); + let engine = engine(&root); + assert_eq!( + ids( + &engine, + "how do I bake sourdough bread", + &vec_of(&[(0, 1.0)]), + &global + ), + Vec::::new() + ); + std::fs::remove_dir_all(&root).ok(); + std::fs::remove_dir_all(&global).ok(); + } + + /// Found by words alone while the query's meaning found nothing near it, + /// a doc needs half the query's weight; with no embedding, a third. + #[test] + fn words_alone_must_say_more_when_meaning_disagrees() { + let (root, global) = (tmp("uncorroborated"), tmp("uncorroborated-global")); + let engine = engine(&root); + // Orthogonal to every doc: the dense leg finds nothing. + let elsewhere = vec_of(&[(100, 1.0)]); + let query = "grep large index"; // doc 7 carries one term of three + assert_eq!( + ids(&engine, query, &elsewhere, &global), + Vec::::new() + ); + let unembedded: Vec = engine + .retrieve_with_vector(query, None, 5, &global) + .into_iter() + .map(|d| d.id) + .collect(); + assert_eq!(unembedded, vec!["doc:7".to_string()]); + // Two terms of four is half: enough on words alone. + assert_eq!( + ids(&engine, "grep prefilter large index", &elsewhere, &global), + vec!["doc:7".to_string()] + ); + std::fs::remove_dir_all(&root).ok(); + std::fs::remove_dir_all(&global).ok(); + } + + /// Sharing only a word every doc uses is not evidence; naming the rare + /// term is. + #[test] + fn a_common_word_alone_does_not_make_a_match() { + let (root, global) = (tmp("common"), tmp("common-global")); + let engine = engine(&root); + assert_eq!( + ids( + &engine, + "grep prefilter large repositories", + &vec_of(&[(0, 1.0)]), + &global + ), + vec!["doc:7".to_string()] + ); + std::fs::remove_dir_all(&root).ok(); + std::fs::remove_dir_all(&global).ok(); + } + + /// Found by meaning alone, a doc that stands out of the corpus is still an + /// answer — and only that one, not the four that merely clear the floor. + #[test] + fn a_doc_that_stands_out_by_meaning_is_kept_alone() { + let (root, global) = (tmp("outlier"), tmp("outlier-global")); + let engine = engine(&root); + assert_eq!( + ids( + &engine, + "something phrased differently", + &vec_of(&[(0, 1.0), (5, 2.0)]), + &global + ), + vec!["doc:5".to_string()] + ); + std::fs::remove_dir_all(&root).ok(); + std::fs::remove_dir_all(&global).ok(); + } +} diff --git a/crates/atlas-native-agent/src/engine/connection.rs b/crates/atlas-native-agent/src/engine/connection.rs index 3de1ce940..411a143d8 100644 --- a/crates/atlas-native-agent/src/engine/connection.rs +++ b/crates/atlas-native-agent/src/engine/connection.rs @@ -721,6 +721,9 @@ impl EngineConnection { // And, for the same reason, organisation access: it may be // handed the organisation tool server (ADR-0014). org_access: true, + // The engine runs in this process and reads the entries from + // memory: a header never reaches a command line (ADR-0020). + in_process: true, cwd: cwd.to_path_buf(), session_id: session_id.cloned(), }, diff --git a/docs/adr/0010-shared-memory-is-pulled-through-the-tool-server.md b/docs/adr/0010-shared-memory-is-pulled-through-the-tool-server.md index 6aef4707e..1d7420f60 100644 --- a/docs/adr/0010-shared-memory-is-pulled-through-the-tool-server.md +++ b/docs/adr/0010-shared-memory-is-pulled-through-the-tool-server.md @@ -1,6 +1,6 @@ # ADR-0010: Shared memory reaches an agent only through the memory tool server -**Status:** Accepted (2026-09-21). Branch `feat/shared-memory-unification`, after spec #77's thirteen tickets landed. +**Status:** Accepted (2026-09-21). Branch `feat/shared-memory-unification`, after spec #77's thirteen tickets landed. Amended by ADR-0020: an ACP agent gets the server through the stdio bridge whatever it advertises; only the in-process native agent gets it over HTTP. **For agents:** this supersedes the "Push" row of `docs/research/shared-memory-system.md` § "Outcome in one page" and rounds 2–3 decisions Q14, Q4 (reopened) and Q15 there. Nothing is prepended to a prompt any more. If a regression is traced to memory not reaching an agent, look at the tool server (`src-tauri/src/commands/memory_server/`) and its instructions, not at `agents_send`. diff --git a/docs/adr/0015-code-search-is-served-by-the-atlas-code-tool-server.md b/docs/adr/0015-code-search-is-served-by-the-atlas-code-tool-server.md index 56bcb5642..29c174b5f 100644 --- a/docs/adr/0015-code-search-is-served-by-the-atlas-code-tool-server.md +++ b/docs/adr/0015-code-search-is-served-by-the-atlas-code-tool-server.md @@ -1,6 +1,6 @@ # ADR-0015: Code search is served by the `atlas_code` tool server -**Status:** Accepted (2026-10-03). Phase 1 of `docs/superpowers/plans/2026-10-03-codeindex-search/`. +**Status:** Accepted (2026-10-03). Phase 1 of `docs/superpowers/plans/2026-10-03-codeindex-search/`. Amended by ADR-0020: an ACP agent gets the server through the stdio bridge whatever it advertises; only the in-process native agent gets it over HTTP. **For agents:** if an agent cannot search code, or searches the wrong place, look at the code tool server (`src-tauri/src/commands/code_server/`) and the engine it calls (`crates/atlas-search`), diff --git a/docs/adr/0019-a-stdio-bridge-carries-memory-to-agents-without-http-mcp.md b/docs/adr/0019-a-stdio-bridge-carries-memory-to-agents-without-http-mcp.md index f702ef270..d24e03433 100644 --- a/docs/adr/0019-a-stdio-bridge-carries-memory-to-agents-without-http-mcp.md +++ b/docs/adr/0019-a-stdio-bridge-carries-memory-to-agents-without-http-mcp.md @@ -1,6 +1,6 @@ # ADR-0019: A stdio bridge carries memory to agents without HTTP MCP -**Status:** Accepted (2026-10-07). Milestone M4 of `docs/superpowers/plans/2026-10-03-memory-system/`. Amends ADR-0010's "an ACP agent that does not advertise HTTP MCP gets no memory at all". +**Status:** Accepted (2026-10-07). Milestone M4 of `docs/superpowers/plans/2026-10-03-memory-system/`. Amends ADR-0010's "an ACP agent that does not advertise HTTP MCP gets no memory at all". Amended by ADR-0020: the bridge now carries every ACP agent, and the token moved from `ATLAS_MCP_TOKEN` to a private file, because an adapter may put a stdio entry's `env` on its own command line. **For agents:** the bridge is `src-tauri/src/commands/memory_bridge.rs`; the offer is `OfferDecision::IncludedViaBridge` in `src-tauri/src/commands/memory_server/offers.rs`; the dispatch is `run_bridge_if_asked` in `src-tauri/src/lib.rs`, called first in `main.rs`. diff --git a/docs/adr/0020-no-credential-rides-an-mcp-entry-handed-to-another-process.md b/docs/adr/0020-no-credential-rides-an-mcp-entry-handed-to-another-process.md new file mode 100644 index 000000000..de7d301bc --- /dev/null +++ b/docs/adr/0020-no-credential-rides-an-mcp-entry-handed-to-another-process.md @@ -0,0 +1,35 @@ +# ADR-0020: No credential rides an MCP entry handed to another process + +**Status:** Accepted (2026-10-09). Amends ADR-0019 (the bridge now carries every ACP agent, and its token moves from `ATLAS_MCP_TOKEN` to a file) and ADR-0010 / ADR-0015 (an ACP agent that advertises HTTP MCP no longer gets the servers over HTTP). + +**For agents:** the rule is `SessionMcpRequest::in_process` in `crates/atlas-agent-servers/src/session_mcp.rs`; the offer is `MemorySessionOffers::offer` in `src-tauri/src/commands/memory_server/offers.rs`; the token files are `TokenFiles` in `src-tauri/src/commands/memory_bridge.rs`; the guarding test is `no_entry_offered_to_another_process_carries_the_token` in `src-tauri/src/commands/memory_server/tests.rs`. + +## Context + +- Atlas offered `atlas_memory` and `atlas_code` to an ACP agent that advertised `mcpCapabilities.http` as `McpServer::Http`, with the session's bearer token in an `Authorization` header. +- The Claude Agent SDK, under the claude-acp adapter, passes every MCP server it is given to the `claude` CLI as `--mcp-config `, so the token sat in the CLI's argv. Any local user can read argv with `ps`. This was seen live with `ps aux | grep -- --mcp-config`. +- The same SDK writes a stdio server's `env` into that JSON as well (the adapter, `@agentclientprotocol/claude-agent-acp`'s `acp-agent.js`, copies an ACP stdio entry's `env` into the SDK's server config; the SDK serialises the whole map into `--mcp-config`). ADR-0019's bridge kept the token out of the bridge's own argv, but the token still reached the adapter's command line through `env`. Putting the token in a different field does not help. +- What an agent does with its MCP configuration is up to the agent: a command line, a config file, a log. Atlas cannot know, and per ADR-0002 must not guess by agent id. + +## Decision + +- **Only an agent in the Atlas process is handed a credential in an MCP entry.** `SessionMcpRequest` carries `in_process`, a property of the connection like `ui_control`: the native connection sets it, an ACP connection never does. The native agent keeps HTTP, with the token in the header its engine reads from memory. +- **Every ACP agent gets the servers through the stdio bridge, whatever it advertises.** ACP requires every agent to accept stdio, so no capability is assumed. The entry is ` mcp-bridge `, with an empty `env`. Every field of it is safe to print. +- **The token lives in a file only the user can read.** Each offer writes its token to a new `0600` file in a `0700` directory of the launch's own, under `/mcp-bridge/`. Those modes are Unix-only: on Windows the file has no mode of its own and inherits the ACL of the per-user `%APPDATA%` tree it sits in, which only that user (and administrators) can read. The bridge reads the file when it starts, so an agent that relaunches the server reads it again. An offer released unbound deletes its file, and a bound one's goes when its session ends. Each write sweeps this launch's files whose token has been revoked. At startup, off the async runtime, Atlas deletes the directories of launches that have ended and were untouched for seven days (`STALE_LAUNCH`), never a live one: a running launch holds a lock on its directory's `.launch.lock` and refreshes its mtime on every write, because two Atlas processes can run on one data directory (the installed app beside a debug build compiled without `tauri.dev.conf.json`, whose identifier is the release one; the single-instance plugin is release-only). A leftover file is harmless, because tokens live only in Atlas's memory and die with the session or the launch. +- Without the Atlas binary's path (`current_exe` fails), an ACP session gets no Atlas servers and the offer's log line says why. If the token file cannot be written, the session gets none either: the token is revoked unused and a warning is logged. Neither case falls back to HTTP. +- The rule covers every server on the shared token, not only memory and code: an entry goes over HTTP only when the request is `in_process` (and the agent takes HTTP); anything else is bridged. The UI and organisation servers are offered only to a connection carrying `ui_control` / `org_access`, which today only the native one does, so in practice they never ride the bridge; if a future connection carried either without being in process, its entries would be bridged too, not sent with a header. +- The bridge's path is `current_exe()`, which is whichever binary is serving: `target/debug/atlas` under `bun run dev:app`, `Atlas.app/Contents/MacOS/atlas` in a release. `main` runs `run_bridge_if_asked` before anything else, so the bridge never boots a window, a profile or the single-instance plugin. +- `preapproval_meta` pre-approves the host's stdio entries for the Claude Code adapter, as it did its HTTP ones. + +## Considered + +- **Keep HTTP for ACP agents that advertise it, and move the token elsewhere.** No other field is any safer: the URL, the headers and `env` all land on the command line. A token-less URL with a one-time ticket in argv can be redeemed by whoever reads `ps` first. HTTP would have saved one local process and one loopback hop per call, which costs nothing anyone would notice. +- **The token in the bridge's `env`** (ADR-0019 as written). It is safe from the bridge's own `ps` line but not from the adapter's, as shown above. +- **Choosing per agent** (HTTP for an adapter known to keep its config in memory). Against ADR-0002, and unknowable: an adapter's next release may change it. + +## Consequences + +- An ACP agent runs one bridge process per offered server per session. +- The log line reads `memory_server=included_via_bridge` for every ACP agent, including one that logs `http_mcp=true`. +- The user's own processes can read the token file. That is the same exposure as `env` (same-user processes can read each other's environment and memory), and the file is closed to other users. +- Windows: the file sits under the user's profile, whose ACL already restricts it. ADR-0019's manual check of the GUI-subsystem bridge launch still applies, and now covers every ACP agent on Windows rather than only stdio-only ones. diff --git a/docs/architecture.md b/docs/architecture.md index 170b66b68..d3f15d92d 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -99,7 +99,7 @@ One Rust module per IPC domain under `src-tauri/src/commands/`. `commands/mod.rs | Domain group | Modules | |---|---| | Agents (ported ACP stack) | agents, agent_host, agent_transcript, agent_analytics, agent_memory, catalog, registry, capture, artifacts_cloud | -| Agent tool servers | memory_server (ADR-0010), memory_bridge (ADR-0019), ui_server (ADR-0012) | +| Agent tool servers | memory_server (ADR-0010), memory_bridge (ADR-0019, ADR-0020), ui_server (ADR-0012) | | Terminal / browser / fs | terminal, browser, fs | | Git | git, git_graph, git_watcher, git_autofetch, gitdiff, git_ops, git_conflicts, git_snapshot, git_stage_ops | | GitHub | github | @@ -184,7 +184,7 @@ The `SessionDelta` shapes those consumers pattern-match live in **`crates/atlas- ### Atlas's tool servers -Atlas hands agents three in-process MCP services on one loopback listener, behind one bearer token per session: the **memory tool server** (`memory_server/`, `/mcp`, ADR-0010), offered to every agent — over HTTP to one that advertises HTTP MCP, else as a stdio server, the Atlas binary run as `atlas mcp-bridge ` with the token in `ATLAS_MCP_TOKEN` (ADR-0019) — the **UI tool server** (`ui_server/`, `/ui`, ADR-0012), offered only to a connection that carries **UI control** — today the in-process native connection — and the **organisation tool server** (below). One offer (`MemorySessionOffers`) decides all three, because the token table holds one token per session. A UI tool call crosses to the window as `atlas:ui-action`; the frontend performs it through the app's own openers (`src/features/ui-actions/`) and answers through `ui_action_respond`, so Rust mirrors no layout or focus state. `tests/ui-actions-contract.test.ts` keeps the tool list and the window's dispatcher in step. +Atlas hands agents three in-process MCP services on one loopback listener, behind one bearer token per session: the **memory tool server** (`memory_server/`, `/mcp`, ADR-0010), offered to every agent — over HTTP to the in-process native agent, and to every ACP agent, whatever it advertises, as a stdio server: the Atlas binary run as `atlas mcp-bridge `, the token in a private file and never in the entry, which an adapter may put on its command line (ADR-0019, ADR-0020) — the **UI tool server** (`ui_server/`, `/ui`, ADR-0012), offered only to a connection that carries **UI control** — today the in-process native connection — and the **organisation tool server** (below). One offer (`MemorySessionOffers`) decides all three, because the token table holds one token per session. A UI tool call crosses to the window as `atlas:ui-action`; the frontend performs it through the app's own openers (`src/features/ui-actions/`) and answers through `ui_action_respond`, so Rust mirrors no layout or focus state. `tests/ui-actions-contract.test.ts` keeps the tool list and the window's dispatcher in step. ### The organisation tool server @@ -275,6 +275,8 @@ Everything else is per-project files under `/.atlas/`: ├── canvas.json ReactFlow node/edge state (Canvas / Spaces) ├── editor-state.json open tabs + split-column layout ├── project.json per-project settings +├── memory-sharing.json shared-memory switches: enabled (default true), fromExternalSessions (default false) +├── memory-summarizer.json handoff summarizer / extraction model choice ├── recent-files.json recent-file list └── git-status-cache.json cached git status ``` diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index be7e479ea..dadc38894 100644 --- a/docs/reference/configuration.md +++ b/docs/reference/configuration.md @@ -542,6 +542,19 @@ schema accepts only `"agent"` or `"off"`. - **Session history, transcripts, per-project `.atlas/` state.** File-backed, but not a "setting," and out of scope until an explicit ownership design says otherwise. +- **Per-project shared-memory switches.** Consent to share memory is given + per repository, so it lives in `/.atlas/memory-sharing.json`, set + from the Memory panel (the Shared pill and its settings popover), not here: + - `enabled` (default `true`): the project's sessions get the memory tools, + their activity is captured, and the extractor runs. + - `fromExternalSessions` (default `false`): extraction also reads the + coding-agent sessions the capture recorder imported from outside Atlas + (a terminal Claude Code session's own transcript), which sends their + transcripts to the extraction model (the Atlas model, or your provider + key), the same as in-app sessions. Has no effect while `enabled` is + `false`. Off, those sessions are never read for extraction; the + briefing's handoff note still names their files and commits, read + locally. - **Any "is a credential configured" presence flag.** Even a boolean is security-sensitive derived state; the BYOK UI computes availability at runtime from the environment instead of persisting it here. diff --git a/src-tauri/src/commands/agents.rs b/src-tauri/src/commands/agents.rs index 8998d32a2..9dccbdb91 100644 --- a/src-tauri/src/commands/agents.rs +++ b/src-tauri/src/commands/agents.rs @@ -614,7 +614,11 @@ enum LifecycleWrite { } impl SharingGatedLifecycle { - fn new(app: AppHandle, server: Arc) -> Self { + fn new( + app: AppHandle, + server: Arc, + token_files: Arc, + ) -> Self { let keep_awake = app .try_state::>() .map(|state| state.inner().clone()); @@ -631,6 +635,8 @@ impl SharingGatedLifecycle { } } LifecycleWrite::Ended { session_id } => { + // Its bridge's token file, dead with its token. + token_files.session_ended(&session_id); // The end-of-session extraction, queued behind the // session's last turn-finished pass. if let Some(ended) = memory.session_ended(&session_id) { @@ -780,9 +786,21 @@ pub fn install_manager(app: &AppHandle) { ))); let server = Arc::new(super::memory_server::MemoryServerHost::new()); app.manage(server.clone()); + // An ACP agent's bridge reads its token from a private file here: + // never from its entry, which may land on a command line (ADR-0020). + // A bound session's file goes when the session ends; ended launches' + // directories are swept once, off the runtime. + let token_files = Arc::new(super::memory_bridge::TokenFiles::new( + data_dir.join("mcp-bridge"), + )); + { + let token_files = token_files.clone(); + tauri::async_runtime::spawn_blocking(move || token_files.sweep_stale_launches()); + } host.set_session_lifecycle(Arc::new(SharingGatedLifecycle::new( app.clone(), server.clone(), + token_files.clone(), ))); let gate_app = app.clone(); let gate: super::memory_server::SharingGate = @@ -878,7 +896,8 @@ pub fn install_manager(app: &AppHandle) { super::org_server::OrgOffer::new(org_access, session_orgs) .describing_with(org_tools), ) - .with_code(super::code_server::CodeOffer::new(code_tools)), + .with_code(super::code_server::CodeOffer::new(code_tools)) + .with_token_files(token_files), )); // `memory_search` also answers from the project's indexed documents. let index_app = app.clone(); diff --git a/src-tauri/src/commands/code_server/offers.rs b/src-tauri/src/commands/code_server/offers.rs index d7a3b0dd1..a1f21051a 100644 --- a/src-tauri/src/commands/code_server/offers.rs +++ b/src-tauri/src/commands/code_server/offers.rs @@ -17,12 +17,13 @@ pub enum CodeOfferDecision { } impl CodeOfferDecision { - /// Included for any agent that advertised HTTP MCP, while the user lets + /// Included for any agent a transport reaches (`reachable`: HTTP in + /// process, the stdio bridge otherwise; ADR-0020), while the user lets /// agents use the code tools and the server is running. Not decided by /// the shared-memory toggle, nor by which agent it is. - pub fn decide(http_mcp: bool, enabled: bool, server_running: bool) -> Self { - if !http_mcp { - Self::Omitted("agent did not advertise mcpCapabilities.http") + pub fn decide(reachable: bool, enabled: bool, server_running: bool) -> Self { + if !reachable { + Self::Omitted("no transport reaches the agent") } else if !enabled { Self::Omitted("code tools are off in Settings") } else if !server_running { @@ -54,8 +55,8 @@ impl CodeOffer { } /// Decide for one request; the setting is read only when it can matter. - pub fn decide(&self, http_mcp: bool, server_running: bool) -> CodeOfferDecision { - let enabled = http_mcp && (self.gate)(); - CodeOfferDecision::decide(http_mcp, enabled, server_running) + pub fn decide(&self, reachable: bool, server_running: bool) -> CodeOfferDecision { + let enabled = reachable && (self.gate)(); + CodeOfferDecision::decide(reachable, enabled, server_running) } } diff --git a/src-tauri/src/commands/code_server/tests.rs b/src-tauri/src/commands/code_server/tests.rs index fa7041db1..98264ef80 100644 --- a/src-tauri/src/commands/code_server/tests.rs +++ b/src-tauri/src/commands/code_server/tests.rs @@ -324,6 +324,7 @@ fn session_request(http_mcp: bool) -> SessionMcpRequest { http_mcp, ui_control: false, org_access: false, + in_process: false, cwd: std::path::PathBuf::from("/p"), session_id: None, } @@ -331,23 +332,7 @@ fn session_request(http_mcp: bool) -> SessionMcpRequest { /// Every entry an offer carries, as `(name, url, bearer token)`. fn entries(offer: &SessionMcpOffer) -> Vec<(String, String, String)> { - offer - .servers() - .iter() - .map(|server| { - let acp::McpServer::Http(http) = server else { - panic!("HTTP entries only") - }; - let token = http - .headers - .iter() - .find(|h| h.name == "Authorization") - .and_then(|h| h.value.strip_prefix("Bearer ")) - .expect("a bearer token") - .to_string(); - (http.name.clone(), http.url.clone(), token) - }) - .collect() + crate::commands::memory_server::offers::offered_entries(offer) } fn names(offer: &SessionMcpOffer) -> Vec { @@ -399,33 +384,28 @@ async fn with_code_tools_off_the_code_server_is_not_offered() { ); } -/// ADR-0019: an agent without HTTP MCP gets the code server through the -/// same stdio bridge as memory, on the same token. +/// ADR-0019, ADR-0020: an ACP agent gets the code server through the same +/// stdio bridge as memory, on the same token, whether or not it advertised +/// HTTP MCP. #[tokio::test(flavor = "multi_thread")] -async fn an_agent_without_http_mcp_gets_the_code_server_through_the_bridge() { +async fn an_acp_agent_gets_the_code_server_through_the_bridge() { let host = running_host(enabled(true)).await; let offers = MemorySessionOffers::new(host, sharing(true)).with_code(CodeOffer::new(enabled(true))); - let offer = offers.offer(&session_request(false)); - let bridged: Vec<(String, String)> = offer - .servers() - .iter() - .map(|server| { - let acp::McpServer::Stdio(stdio) = server else { - panic!("stdio entries only: {server:?}") - }; - let token = stdio - .env + for http_mcp in [true, false] { + let offer = offers.offer(&session_request(http_mcp)); + assert!( + offer + .servers() .iter() - .find(|v| v.name == "ATLAS_MCP_TOKEN") - .map(|v| v.value.clone()) - .expect("the token rides the environment"); - (stdio.name.clone(), token) - }) - .collect(); - let names: Vec<&str> = bridged.iter().map(|(n, _)| n.as_str()).collect(); - assert_eq!(names, ["atlas_memory", "atlas_code"]); - assert_eq!(bridged[0].1, bridged[1].1, "one token for both"); + .all(|s| matches!(s, acp::McpServer::Stdio(_))), + "stdio entries only: {:?}", + offer.servers() + ); + let got = entries(&offer); + assert_eq!(names(&offer), ["atlas_memory", "atlas_code"]); + assert_eq!(got[0].2, got[1].2, "one token for both"); + } } #[test] @@ -434,7 +414,7 @@ fn the_decision_says_whether_the_code_server_is_included_and_why_not() { assert_eq!(CodeOfferDecision::decide(true, true, true), Included); assert_eq!( CodeOfferDecision::decide(false, true, true), - Omitted("agent did not advertise mcpCapabilities.http") + Omitted("no transport reaches the agent") ); assert_eq!( CodeOfferDecision::decide(true, false, true), diff --git a/src-tauri/src/commands/memory_bridge.rs b/src-tauri/src/commands/memory_bridge.rs index c919da845..5ffbe1225 100644 --- a/src-tauri/src/commands/memory_bridge.rs +++ b/src-tauri/src/commands/memory_bridge.rs @@ -1,7 +1,10 @@ -//! `atlas mcp-bridge `: a stdio MCP server that forwards every message -//! to Atlas's own loopback tool server (ADR-0019), for agents that can only -//! launch stdio MCP servers. The session token arrives in `ATLAS_MCP_TOKEN`, -//! never on the command line. +//! `atlas mcp-bridge `: a stdio MCP server that forwards +//! every message to Atlas's own loopback tool server (ADR-0019), for every +//! agent that runs in another process (ADR-0020). The session token is read +//! from a file only this user can read ([`TokenFiles`]); the command line +//! names the file, never the token, because an agent may put the whole MCP +//! entry — arguments and environment alike — on its own command line, where +//! any local user's `ps` reads it. //! //! One JSON-RPC message per input line is POSTed to the server; every //! message in the answer (a JSON body, or each SSE `data:` line) is written @@ -14,6 +17,243 @@ use anyhow::{bail, Result}; use serde_json::Value; use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncWrite, AsyncWriteExt}; +/// Where offered sessions' tokens are written for the bridge to read: one +/// file per offer, readable only by this user, in a directory of this +/// launch's own that only this user can enter. +/// +/// Tokens live in Atlas's memory, so a file whose token is dead holds nothing +/// that opens anything; the files are tidied all the same. An offer released +/// unbound removes its file, and a bound one's goes when its session ends +/// ([`session_ended`](Self::session_ended)); every write sweeps this launch's +/// files whose token has since been revoked; and at startup +/// ([`sweep_stale_launches`](Self::sweep_stale_launches), off the runtime) +/// the directories of launches that have ended and were untouched for +/// [`STALE_LAUNCH`] go. Never another live launch's directory: a dev build +/// and the installed app run side by side on one data directory, and a +/// bridge an agent relaunches reads its file again. A live launch holds a +/// lock on its directory's [`LOCK_FILE`] for as long as it runs, and +/// refreshes that file's mtime on every write. +pub struct TokenFiles { + root: std::path::PathBuf, + dir: std::path::PathBuf, + /// This launch's directory, made on the first write; holds the lock. + prepared: std::sync::OnceLock>, + written: std::sync::Mutex>, + /// A test's root, removed when it drops (even after a failed assertion). + #[cfg(test)] + _temp: Option, +} + +/// One file [`TokenFiles::write`] made: its token, and the session it was +/// bound to once the agent answered. +struct Written { + path: std::path::PathBuf, + token: String, + session: Option, +} + +/// A launch directory not written for this long, whose launch has ended, +/// is swept. +pub const STALE_LAUNCH: std::time::Duration = std::time::Duration::from_secs(7 * 24 * 60 * 60); + +/// The file in a launch directory its launch holds locked while it runs. +pub const LOCK_FILE: &str = ".launch.lock"; + +impl TokenFiles { + /// Token files under `root`, which belongs to Atlas alone, in a directory + /// of this launch's own. + pub fn new(root: std::path::PathBuf) -> Self { + let dir = root.join(uuid::Uuid::new_v4().simple().to_string()); + Self { + root, + dir, + prepared: std::sync::OnceLock::new(), + written: std::sync::Mutex::new(Vec::new()), + #[cfg(test)] + _temp: None, + } + } + + /// Token files under a fresh directory of their own in the system's + /// temporary directory: for a host that names none. Nothing is created + /// until the first write. + #[cfg(not(test))] + pub fn in_temp() -> Self { + Self::new(std::env::temp_dir().join(format!( + "atlas-mcp-bridge-{}", + uuid::Uuid::new_v4().simple() + ))) + } + + /// In tests, a [`tempfile::TempDir`] removed when these files drop. + #[cfg(test)] + pub fn in_temp() -> Self { + let temp = tempfile::Builder::new() + .prefix("atlas-mcp-bridge-test-") + .tempdir() + .expect("a temporary directory"); + let mut files = Self::new(temp.path().join("mcp-bridge")); + files._temp = Some(temp); + files + } + + fn written(&self) -> std::sync::MutexGuard<'_, Vec> { + self.written + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } + + /// Write `token` to a new file of its own and return its path, first + /// removing this launch's files whose token `is_live` no longer admits. + /// A few small syscalls: the stale-launch sweep is + /// [`sweep_stale_launches`](Self::sweep_stale_launches)'s, at startup. + pub fn write(&self, token: &str, is_live: impl Fn(&str) -> bool) -> Result { + let lock = self + .prepared + .get_or_init(|| self.prepare().map_err(|e| format!("{e:#}"))) + .as_ref() + .map_err(|e| anyhow::Error::msg(e.clone()))?; + // This launch is in use: its directory reads as fresh to a sweep. + let _ = lock.set_modified(std::time::SystemTime::now()); + let mut written = self.written(); + written.retain(|w| { + let live = is_live(&w.token); + if !live { + let _ = std::fs::remove_file(&w.path); + } + live + }); + let path = self + .dir + .join(format!("{}.token", uuid::Uuid::new_v4().simple())); + let mut options = std::fs::OpenOptions::new(); + options.write(true).create_new(true); + #[cfg(unix)] + std::os::unix::fs::OpenOptionsExt::mode(&mut options, 0o600); + let mut file = options.open(&path)?; + std::io::Write::write_all(&mut file, token.as_bytes())?; + written.push(Written { + path: path.clone(), + token: token.to_string(), + session: None, + }); + Ok(path) + } + + /// The file at `path` now belongs to `session_id`: it goes when that + /// session ends. Any earlier file of that session goes now (binding a + /// new token revoked the old one). + pub fn bind(&self, path: &std::path::Path, session_id: &str) { + let mut written = self.written(); + written.retain(|w| { + let earlier = w.path != path && w.session.as_deref() == Some(session_id); + if earlier { + let _ = std::fs::remove_file(&w.path); + } + !earlier + }); + if let Some(w) = written.iter_mut().find(|w| w.path == path) { + w.session = Some(session_id.to_string()); + } + } + + /// Remove the files of `session_id`, which has ended. Idempotent. + pub fn session_ended(&self, session_id: &str) { + self.written().retain(|w| { + let ended = w.session.as_deref() == Some(session_id); + if ended { + let _ = std::fs::remove_file(&w.path); + } + !ended + }); + } + + /// Remove a file [`write`](Self::write) made. Idempotent. + pub fn remove(&self, path: &std::path::Path) { + self.written().retain(|w| w.path != path); + if path.parent() == Some(self.dir.as_path()) { + let _ = std::fs::remove_file(path); + } + } + + /// Remove the directories of ended launches untouched for + /// [`STALE_LAUNCH`]: never this launch's, nor one whose [`LOCK_FILE`] is + /// held (its launch is running, however long since it wrote). Blocking + /// file I/O, run once at startup on a blocking thread. + pub fn sweep_stale_launches(&self) { + let Ok(entries) = std::fs::read_dir(&self.root) else { + return; + }; + for entry in entries.flatten() { + let path = entry.path(); + if path == self.dir || !entry.file_type().is_ok_and(|t| t.is_dir()) { + continue; + } + if launch_is_stale(&path) { + let _ = std::fs::remove_dir_all(&path); + } + } + } + + /// The root, enterable only by this user; this launch's directory made + /// with `create_dir`, never through an existing path, so nobody else's + /// directory (or a link to one) is adopted; and its lock, held from now + /// until this launch ends. + fn prepare(&self) -> Result { + std::fs::create_dir_all(&self.root)?; + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions(&self.root, std::fs::Permissions::from_mode(0o700))?; + } + let mut builder = std::fs::DirBuilder::new(); + #[cfg(unix)] + std::os::unix::fs::DirBuilderExt::mode(&mut builder, 0o700); + builder.create(&self.dir)?; + let lock = std::fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(self.dir.join(LOCK_FILE))?; + lock.try_lock() + .map_err(|e| anyhow::anyhow!("cannot lock this launch's token directory: {e}"))?; + Ok(lock) + } +} + +/// Whether the launch directory `dir` belongs to a launch that has ended and +/// was untouched for [`STALE_LAUNCH`]: by its [`LOCK_FILE`]'s mtime (which +/// every write refreshes) when it has one, else by the directory's own. +fn launch_is_stale(dir: &std::path::Path) -> bool { + let lock_path = dir.join(LOCK_FILE); + let modified = std::fs::metadata(&lock_path) + .or_else(|_| std::fs::metadata(dir)) + .and_then(|m| m.modified()); + let old = modified + .ok() + .and_then(|t| t.elapsed().ok()) + .is_some_and(|age| age > STALE_LAUNCH); + if !old { + return false; + } + // A held lock is a running launch. A lock file that cannot be opened is + // left alone: nothing proves its launch has ended. + match std::fs::OpenOptions::new().write(true).open(&lock_path) { + Ok(file) => file.try_lock().is_ok(), + Err(e) => e.kind() == std::io::ErrorKind::NotFound, + } +} + +/// The token in the file the bridge was pointed at. +pub fn read_token_file(path: &std::path::Path) -> Result { + let token = std::fs::read_to_string(path) + .map_err(|e| anyhow::anyhow!("cannot read the token file {}: {e}", path.display()))?; + let token = token.trim(); + if token.is_empty() { + bail!("the token file {} is empty", path.display()); + } + Ok(token.to_string()) +} + /// Forward `input` to the server at `url` until `input` ends. pub async fn bridge( url: &str, @@ -306,6 +546,126 @@ mod tests { format!("http://{addr}/mcp") } + /// The token travels in a file only this user can read, in a directory + /// only this user can enter; the bridge reads it back. + #[test] + fn a_token_file_is_private_and_reads_back() { + let files = TokenFiles::in_temp(); + let path = files.write("secret-token", |_| true).unwrap(); + assert_eq!(read_token_file(&path).unwrap(), "secret-token"); + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + let mode = |p: &std::path::Path| std::fs::metadata(p).unwrap().permissions().mode(); + assert_eq!(mode(&path) & 0o777, 0o600); + assert_eq!(mode(path.parent().unwrap()) & 0o777, 0o700); + assert_eq!(mode(&files.root) & 0o777, 0o700); + } + files.remove(&path); + assert!(read_token_file(&path).is_err()); + } + + /// A file whose token was revoked goes at the next write; another live + /// launch's directory under the same root is left alone. + #[test] + fn dead_tokens_files_are_swept_and_other_launches_are_left_alone() { + let temp = tempfile::TempDir::new().unwrap(); + let root = temp.path().to_path_buf(); + let other = TokenFiles::new(root.clone()); + let theirs = other.write("theirs", |_| true).unwrap(); + let files = TokenFiles::new(root); + let dead = files.write("dead", |_| true).unwrap(); + let live = files.write("live", |_| true).unwrap(); + let next = files.write("next", |t| t != "dead").unwrap(); + assert!(!dead.exists(), "a revoked token's file is swept"); + assert!(live.exists() && next.exists()); + files.sweep_stale_launches(); + assert!(theirs.exists(), "another launch's file stays"); + } + + /// Set `path`'s mtime to longer ago than [`STALE_LAUNCH`]. + fn age(path: &std::path::Path) { + let then = std::time::SystemTime::now() - STALE_LAUNCH - std::time::Duration::from_secs(60); + std::fs::File::options() + .write(true) + .open(path) + .unwrap() + .set_modified(then) + .unwrap(); + } + + /// A running launch's directory is never swept, however long since it + /// last wrote; an ended one's is, once untouched for [`STALE_LAUNCH`]; + /// an ended one written recently is not yet. + #[test] + fn only_ended_launches_untouched_for_a_week_are_swept() { + let temp = tempfile::TempDir::new().unwrap(); + let root = temp.path().to_path_buf(); + let running = TokenFiles::new(root.clone()); + let running_file = running.write("r", |_| true).unwrap(); + age(&running.dir.join(LOCK_FILE)); + let ended = TokenFiles::new(root.clone()); + let ended_file = ended.write("e", |_| true).unwrap(); + let ended_dir = ended.dir.clone(); + drop(ended); + age(&ended_dir.join(LOCK_FILE)); + let recent = TokenFiles::new(root.clone()); + let recent_file = recent.write("n", |_| true).unwrap(); + drop(recent); + + let starting = TokenFiles::new(root); + starting.sweep_stale_launches(); + assert!(running_file.exists(), "a running launch keeps its files"); + assert!( + !ended_file.exists() && !ended_dir.exists(), + "an ended, week-old launch goes" + ); + assert!( + recent_file.exists(), + "an ended launch written recently stays for now" + ); + } + + /// Every write marks the launch as in use, so its age is the time since + /// its last write, not since it started. + #[test] + fn a_write_refreshes_the_launch() { + let files = TokenFiles::in_temp(); + files.write("a", |_| true).unwrap(); + let lock = files.dir.join(LOCK_FILE); + age(&lock); + assert!(launch_is_stale_by_age(&lock)); + files.write("b", |_| true).unwrap(); + assert!(!launch_is_stale_by_age(&lock)); + } + + fn launch_is_stale_by_age(lock: &std::path::Path) -> bool { + std::fs::metadata(lock) + .unwrap() + .modified() + .unwrap() + .elapsed() + .is_ok_and(|a| a > STALE_LAUNCH) + } + + /// A bound file goes when its session ends; one bound to another session + /// stays; rebinding a session drops its earlier file. + #[test] + fn a_bound_file_goes_when_its_session_ends() { + let files = TokenFiles::in_temp(); + let first = files.write("t1", |_| true).unwrap(); + let other = files.write("t2", |_| true).unwrap(); + files.bind(&first, "s1"); + files.bind(&other, "s2"); + let again = files.write("t3", |_| true).unwrap(); + files.bind(&again, "s1"); + assert!(!first.exists(), "the session's earlier file goes on rebind"); + files.session_ended("s1"); + assert!(!again.exists(), "an ended session's file goes"); + assert!(other.exists(), "another session's file stays"); + files.session_ended("s1"); + } + #[test] fn an_answer_is_its_sse_data_lines_or_its_body() { let sse = "id: 0\ndata:\n\nevent: message\ndata: {\"id\":1}\n\n"; diff --git a/src-tauri/src/commands/memory_capture.rs b/src-tauri/src/commands/memory_capture.rs index 276312c99..a2116582b 100644 --- a/src-tauri/src/commands/memory_capture.rs +++ b/src-tauri/src/commands/memory_capture.rs @@ -27,6 +27,27 @@ pub struct CaptureReader { /// One scope's capture stores, opened for reading. pub struct ScopeStores(Vec<(PathBuf, Store)>); +/// Which of the recorder's sessions a read takes. The recorder holds the +/// sessions Atlas hosted (`native`, `acp`) and the ones its importer read off +/// an agent's own transcript (`external_jsonl`): a terminal session Atlas +/// never hosted, or a hosted one live capture missed. Extraction reads with +/// the reach `MemorySharingState::extraction_reach` decides; the local reads +/// (the briefing's handoff, provenance) take everything. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Reach { + /// Only the sessions Atlas recorded live: never an imported one. + Hosted, + /// The imported sessions too. + WithImported, +} + +impl Reach { + /// Whether a session recorded under `source` is within this reach. + pub fn admits(self, source: Source) -> bool { + self == Self::WithImported || source != Source::ExternalJsonl + } +} + /// A recorded session, with the store and launch directory it lives in. pub struct Recorded<'a> { pub root: &'a Path, @@ -55,6 +76,67 @@ impl CaptureReader { } impl ScopeStores { + /// The recorded sessions of every agent whose last activity lies in + /// `since_ms..=idle_before_ms`, newest first, one per session id. This + /// includes the terminal sessions the importer read off disk, which + /// Atlas never hosted: the recorder is the only place their conversation + /// reaches Atlas. A session with no activity stamp falls back to its row + /// clock. + pub fn recent_sessions( + &self, + since_ms: i64, + idle_before_ms: i64, + ) -> Vec { + self.session_activity(Reach::WithImported) + .into_iter() + .filter(|(_, at)| (since_ms..=idle_before_ms).contains(at)) + .map(|(writer, _)| writer) + .collect() + } + + /// Every recorded session within `reach` with a native id, once per id, + /// with its last activity (the newest of its rows), newest first. + pub fn session_activity( + &self, + reach: Reach, + ) -> Vec<(crate::commands::shared_memory::Writer, i64)> { + let mut newest: HashMap = HashMap::new(); + for (root, store) in &self.0 { + for workspace in workspace_ids(root) { + for session in store.sessions_for_project(&workspace).unwrap_or_default() { + if session.native_session_id.is_empty() || !reach.admits(session.source) { + continue; + } + let at = session + .last_activity_at + .unwrap_or(session.updated_at) + .timestamp_millis(); + let agent = session + .agent + .clone() + .filter(|a| !a.is_empty()) + .unwrap_or_else(|| session.source.as_str().to_string()); + let slot = newest + .entry(session.native_session_id.clone()) + .or_insert((at, agent.clone())); + if at > slot.0 { + *slot = (at, agent); + } + } + } + } + let mut rows: Vec<(String, (i64, String))> = newest.into_iter().collect(); + rows.sort_by(|a, b| b.1 .0.cmp(&a.1 .0).then_with(|| a.0.cmp(&b.0))); + rows.into_iter() + .map(|(session_id, (at, agent))| { + ( + crate::commands::shared_memory::Writer { agent, session_id }, + at, + ) + }) + .collect() + } + pub fn is_empty(&self) -> bool { self.0.is_empty() } @@ -65,12 +147,20 @@ impl ScopeStores { /// imported from its own transcript (`external_jsonl`). At most six /// lookups on the unique index per store. pub fn find(&self, session_id: &str) -> Option> { + self.find_within(session_id, Reach::WithImported) + } + + /// [`Self::find`], taking only a session recorded within `reach`. + pub fn find_within(&self, session_id: &str, reach: Reach) -> Option> { if session_id.is_empty() { return None; } for (root, store) in &self.0 { for workspace in workspace_ids(root) { for source in [Source::Native, Source::Acp, Source::ExternalJsonl] { + if !reach.admits(source) { + continue; + } let Ok(Some(row)) = store.session_id_for(&workspace, source, session_id) else { continue; }; diff --git a/src-tauri/src/commands/memory_delta.rs b/src-tauri/src/commands/memory_delta.rs index 5aa8ae174..fbc1b2f7d 100644 --- a/src-tauri/src/commands/memory_delta.rs +++ b/src-tauri/src/commands/memory_delta.rs @@ -62,6 +62,12 @@ pub fn ingest(envelope: &SessionDeltaEnvelope, store: &SharedMemoryStore) { EventKind::Architecture => Some(EntryKind::Architecture), _ => None, }; + // A file outside the project (another checkout, `/tmp`, `~/.claude`) + // is not this project's change: it would be briefed to every agent + // here as if it were. + if ev.kind == EventKind::FileChanged && !in_project(&meta.cwd, &ev.key) { + continue; + } let result = match durable { Some(kind) => { let text = ev @@ -85,6 +91,15 @@ pub fn ingest(envelope: &SessionDeltaEnvelope, store: &SharedMemoryStore) { } } +/// Whether `path` is a file of the project the session at `cwd` works in: +/// under `cwd` or under its scope (the repository's root and worktrees). +fn in_project(cwd: &str, path: &str) -> bool { + if atlas_memory::record::path_in_dirs(path, &[std::path::PathBuf::from(cwd)]) { + return true; + } + super::shared_memory::store_for(cwd).is_ok_and(|store| store.in_scope(path)) +} + /// One background thread that runs capture jobs in the order they were /// pushed. The blocking pool gives no ordering across threads, so two plan /// updates sent there could append out of order and the older plan would @@ -132,12 +147,19 @@ pub fn classify(delta: &SessionDelta, session_id: &str, agent: &str) -> Vec SessionDelta { + SessionDelta::ToolCallUpserted { + message_id: "m1".into(), + tool_call: atlas_agent_wire::ToolCall { + id: "call-x".into(), + tool_name: "Edit".into(), + title: Some("Edit".into()), + kind: Some("edit".into()), + status: ToolCallStatus::Completed, + arguments: serde_json::json!({ "paths": paths }), + result: None, + locations: paths + .iter() + .map(|p| serde_json::json!({ "path": p })) + .collect(), + raw_output: None, + content_blocks: vec![], + }, + } + } + + fn envelope(delta: SessionDelta) -> SessionDeltaEnvelope { + SessionDeltaEnvelope { + agent_id: atlas_agent_wire::AgentId::new(), + session_id: "s1".into(), + delta, + } + } + + /// A plan whose every item is completed is finished: it is logged as + /// `done`, which clears the active plan, rather than lingering as active. + #[test] + fn a_finished_plan_clears_the_active_plan() { + let finished = SessionDelta::PlanUpdated { + plan: vec![ + PlanEntry { + content: "Create branch".into(), + priority: None, + status: "completed".into(), + }, + PlanEntry { + content: "Open PR".into(), + priority: None, + status: "completed".into(), + }, + ], + }; + let evs = classify(&finished, "s1", "claude-acp"); + assert_eq!(evs[0].payload["status"], "done"); + assert_eq!( + classify(&plan_delta(), "s1", "x")[0].payload["status"], + "active" + ); + + let dir = std::env::temp_dir().join(format!("atlas-delta-plan-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let project = dir.to_string_lossy().to_string(); + let store = SharedMemoryStore::new(); + store.register_session("s1", &project, "claude-acp"); + ingest(&envelope(plan_delta()), &store); + assert_eq!(store.list_entries(&project, Some(EntryKind::Plan)).len(), 1); + ingest(&envelope(finished), &store); + assert!(store + .list_entries(&project, Some(EntryKind::Plan)) + .is_empty()); + let _ = std::fs::remove_dir_all(&dir); + } + + /// An edit outside the project (another checkout, `/tmp`, `~/.claude`) + /// is not recorded into this project's memory. + #[test] + fn an_edit_outside_the_project_is_not_captured() { + let dir = std::env::temp_dir().join(format!("atlas-delta-scope-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let project = dir.to_string_lossy().to_string(); + let inside = format!("{project}/src/lib.rs"); + let store = SharedMemoryStore::new(); + store.register_session("s1", &project, "claude-acp"); + ingest( + &envelope(edit_delta(&[ + &inside, + "src/main.rs", + "/tmp/homebrew-cask-pr/Casks/a/atlas-ai.rb", + "/Users/someone/Developer/atlas/src/App.tsx", + "/Users/someone/.claude/plans/x.md", + ])), + &store, + ); + let mut keys: Vec = store + .list_entries(&project, Some(EntryKind::FileChanged)) + .into_iter() + .map(|e| e.key) + .collect(); + keys.sort(); + assert_eq!(keys, vec![inside, "src/main.rs".to_string()]); + let _ = std::fs::remove_dir_all(&dir); + } + #[test] fn decision_marker_extracted() { let evs = scan_assistant_text("We will use RS256 for signing.", "s1", "codex"); diff --git a/src-tauri/src/commands/memory_extract.rs b/src-tauri/src/commands/memory_extract.rs index e107e06d1..a62819b56 100644 --- a/src-tauri/src/commands/memory_extract.rs +++ b/src-tauri/src/commands/memory_extract.rs @@ -169,6 +169,59 @@ impl Extractor { .await } + /// Whether this process holds `session_id`'s turns for its end pass + /// (a turn of it finished here since it last started). + pub fn holds_turns(&self, session_id: &str) -> bool { + self.turns.lock().contains_key(session_id) + } + + /// The end pass of a session whose turns this process does not hold, + /// over its persisted conversation (`turns`, see [`persisted_turns`]). + /// + /// The end pass normally runs over the turns the turn-finished passes + /// kept in memory. Those are gone when the app quits, and quitting is how + /// most sessions end: the end is recorded on the way out and its pass, a + /// model call, is cut off with the process. A session too short for the + /// turn gates then contributes nothing, ever. This pass makes up for it, + /// for an end that ran without turns and for every recent session the + /// health pass finds that is no longer live. The persisted counters make + /// it run once: nothing new since the last pass, no model call. + pub async fn catch_up( + &self, + sharing: &MemorySharingState, + cwd: &str, + writer: &Writer, + turns: Vec, + ) -> usize { + if self.holds_turns(&writer.session_id) { + return 0; // live here: its own end pass covers it + } + let Some(route) = self.route(sharing, cwd) else { + return 0; + }; + // A long session is read in prompt-sized passes from where the last + // pass stopped, not in one pass that would see only its tail; at + // most CATCH_UP_PASSES now, the next job carries on. + let mut recorded = 0; + for _ in 0..CATCH_UP_PASSES { + let (ran, stored) = self + .run_pass( + route.clone(), + cwd, + writer, + &turns, + Trigger::SessionEnd, + Some(CATCH_UP_CHUNK_CHARS), + ) + .await; + recorded += stored; + if !ran { + break; + } + } + recorded + } + /// The model a pass in `cwd` would ask, or `None` when no pass runs there /// (sharing off, the reserved local mode, no account and no BYOK choice). /// The model passes ask (the dream pass asks it too, with the same @@ -194,6 +247,24 @@ impl Extractor { turns: &[TranscriptTurn], trigger: Trigger, ) -> usize { + self.run_pass(route, cwd, writer, turns, trigger, None) + .await + .1 + } + + /// One pass; `(whether it ran, entries recorded)`. With `chunk`, the + /// pass reads only the turns from where the last one stopped up to about + /// `chunk` characters of text ([`chunk_end`]), so the next pass can read + /// on from there. + async fn run_pass( + &self, + route: Route, + cwd: &str, + writer: &Writer, + turns: &[TranscriptTurn], + trigger: Trigger, + chunk: Option, + ) -> (bool, usize) { // The gate counters, from the scope's memory directory (git lookup // and file reads: off the async runtime). let loaded = { @@ -210,10 +281,23 @@ impl Extractor { Ok(loaded) => loaded, Err(e) => { tracing::debug!(target: "atlas::shared_memory", "extraction skipped: {e}"); - return 0; + return (false, 0); } }; let passes_before = state.extraction_count; + // A chunked pass reads turns whose outside reads are known per turn + // (the recorder's), and its prompt holds only the chunk: whether it + // read outside content is the chunk's. An unchunked pass judges the + // whole conversation. + let start = state.last_extracted_turn_index.min(turns.len()); + let (turns, judged) = match chunk { + Some(chars) => { + let end = chunk_end(turns, start, chars); + (&turns[..end], &turns[start..end]) + } + None => (turns, turns), + }; + let external = judged.iter().any(|t| t.external); let model = self.model.clone(); let found = extract::extract(turns, &mut state, trigger, |prompt| async move { @@ -229,21 +313,23 @@ impl Extractor { let found = match found { Ok(found) => found, Err(e) => { - tracing::debug!(target: "atlas::shared_memory", "extraction pass failed: {e:#}"); - return 0; + // Warn, not debug: a pass that keeps failing (the gateway + // unreachable, no model for the org) is otherwise invisible, + // and memory silently stops growing. + tracing::warn!(target: "atlas::shared_memory", "extraction pass failed: {e:#}"); + return (false, 0); } }; if state.extraction_count == passes_before { - return 0; // no pass ran (the gates are not met yet): nothing to save + return (false, 0); // no pass ran (the gates are not met yet): nothing to save } - let external = turns.iter().any(|t| t.external); // Persist the counters and land the entries (SQLite writes: off the // async runtime). let (memory, cwd, writer) = (self.memory.clone(), cwd.to_string(), writer.clone()); - tokio::task::spawn_blocking(move || { + let recorded = tokio::task::spawn_blocking(move || { if let Err(e) = state.save(&memory_dir, &writer.session_id) { - tracing::debug!(target: "atlas::shared_memory", "extraction state not saved: {e:#}"); + tracing::warn!(target: "atlas::shared_memory", "extraction state not saved: {e:#}"); } let mut recorded = 0; for entry in found { @@ -256,13 +342,13 @@ impl Extractor { }; match memory.record_extracted(&cwd, &writer, entry.kind, &entry.content, confidence) { Ok(_) => recorded += 1, - Err(e) => tracing::debug!(target: "atlas::shared_memory", "extracted entry not recorded: {e}"), + Err(e) => tracing::warn!(target: "atlas::shared_memory", "extracted entry not recorded: {e}"), } } recorded }) - .await - .unwrap_or(0) + .await; + (true, recorded.unwrap_or(0)) } } @@ -288,6 +374,252 @@ pub fn transcript_turns(messages: &[atlas_agent_wire::Message]) -> Vec usize { + let start = start.min(turns.len()); + let (mut used, mut answered) = (0usize, false); + for (i, turn) in turns.iter().enumerate().skip(start) { + let len = turn.text.trim().len(); + if used + len > chars && answered { + return i; + } + used += len; + answered |= turn.role == "assistant" && len > 0; + } + turns.len() +} + +/// A session must have been idle this long before a catch-up pass reads it: +/// one still running — in another Atlas process on the same data directory +/// (a dev build beside the installed app), or a terminal agent — is read once +/// it pauses, not mid-run on every scan. +pub const CATCH_UP_IDLE_MS: i64 = 15 * 60 * 1000; + +/// Of `candidates` (each session with its last activity), the ones a +/// catch-up pass looks at: active within [`CATCH_UP_WINDOW_MS`] of `now`, +/// idle for [`CATCH_UP_IDLE_MS`], not `live` in this process, newest first, +/// at most [`CATCH_UP_MAX`]. Whether one has anything new is the extractor's +/// call (its persisted counters). +fn pick_catch_up( + candidates: Vec<(Writer, i64)>, + now: i64, + live: &dyn Fn(&str) -> bool, +) -> Vec { + let window = (now - CATCH_UP_WINDOW_MS)..=(now - CATCH_UP_IDLE_MS); + let mut rows: Vec<_> = candidates + .into_iter() + .filter(|(w, at)| !w.agent.is_empty() && window.contains(at) && !live(&w.session_id)) + .collect(); + rows.sort_by(|a, b| { + b.1.cmp(&a.1) + .then_with(|| a.0.session_id.cmp(&b.0.session_id)) + }); + rows.into_iter() + .take(CATCH_UP_MAX) + .map(|(w, _)| w) + .collect() +} + +/// The sessions of `store` a catch-up pass looks at ([`pick_catch_up`]). A +/// session's last activity is the newest of its start, its end and its +/// events in the log, so one still running in another process is left +/// alone however long ago it started. +pub fn catch_up_sessions( + store: &atlas_memory::record::RecordStore, + now: i64, + live: &dyn Fn(&str) -> bool, +) -> Vec { + let last_event = store + .last_event_by_session(now - CATCH_UP_WINDOW_MS) + .unwrap_or_default(); + let candidates = store + .sessions() + .unwrap_or_default() + .into_iter() + .filter_map(|s| { + let at = [ + s.started_at, + s.ended_at, + last_event.get(&s.session_id).copied(), + ] + .into_iter() + .flatten() + .max()?; + Some(( + Writer { + agent: s.agent, + session_id: s.session_id, + }, + at, + )) + }) + .collect(); + pick_catch_up(candidates, now, live) +} + +/// The recorded sessions (any agent; the imported terminal ones only when +/// `reach` admits them, see `MemorySharingState::extraction_reach`) a +/// catch-up pass looks at ([`pick_catch_up`]), by their recorded activity. +pub fn recorded_catch_up_sessions( + stores: &super::memory_capture::ScopeStores, + reach: super::memory_capture::Reach, + now: i64, + live: &dyn Fn(&str) -> bool, +) -> Vec { + pick_catch_up(stores.session_activity(reach), now, live) +} + +/// The recorded and the hosted candidates as one list, each session once +/// (the recorder's spelling of its agent wins), at most [`CATCH_UP_MAX`]. +pub fn merge_catch_up(recorded: Vec, hosted: Vec) -> Vec { + let mut seen = std::collections::HashSet::new(); + recorded + .into_iter() + .chain(hosted) + .filter(|w| seen.insert(w.session_id.clone())) + .take(CATCH_UP_MAX) + .collect() +} + +/// A session's conversation as Atlas persisted it, for an end pass that runs +/// after the live turns are gone; `None` when Atlas recorded none. +/// +/// The capture recorder is read first ([`recorded_turns`]): it holds every +/// agent's sessions, including the terminal ones its importer read off disk, +/// which Atlas never hosted and so has no other record of. Those are read +/// only when `reach` admits them (`MemorySharingState::extraction_reach`); +/// otherwise such a session has no conversation here at all. Atlas's own +/// transcript (`agent_transcript`) is the fallback, looked up under `cwd`, +/// then in every project directory (the session may have run in a worktree +/// or a subdirectory of the scope). It keeps text, not tool calls, so +/// whether that session read outside content is unknown, and the pass is +/// treated as having read some: its entries are candidates (M0). +pub fn persisted_turns( + config_dir: &std::path::Path, + cwd: &str, + session_id: &str, + reader: &super::memory_capture::CaptureReader, + reach: super::memory_capture::Reach, +) -> Option> { + use super::agent_transcript as transcripts; + let stores = reader.stores(cwd); + let found = stores.find_within(session_id, reach); + if let Some(turns) = found.as_ref().map(recorded_turns).filter(|t| !t.is_empty()) { + return Some(turns); + } + let stored = transcripts::read(config_dir, cwd, session_id).or_else(|| { + let name = format!("{}.json", transcripts::sanitize_id(session_id)); + std::fs::read_dir(config_dir.join("agent-transcripts")) + .ok()? + .flatten() + .find_map(|dir| transcripts::read_file(&dir.path().join(&name))) + })?; + let external = match &found { + Some(found) => found + .store + .tool_calls_for_session(&found.session.id) + .unwrap_or_default() + .iter() + .any(recorded_call_is_external), + None => true, + }; + let mut turns: Vec = stored + .messages + .iter() + .map(|m| TranscriptTurn { + role: m.role.clone(), + text: atlas_agent_transcript::strip_injected_context(&m.content), + tool_calls: 0, + external: false, + }) + .collect(); + for turn in &mut turns { + turn.external = external; + } + Some(turns) +} + +/// A recorded session's conversation as the extractor reads it: its text +/// messages in order (thinking and tool rows left out), Atlas's injected +/// blocks stripped. A turn's tool calls are counted on its last assistant +/// message, and every message of a turn that made an outside call (a fetch, +/// a third-party MCP tool) is marked external. +pub fn recorded_turns(found: &super::memory_capture::Recorded<'_>) -> Vec { + use atlas_checkpoint::{Mode, Role}; + let messages = found + .store + .messages_for_session(&found.session.id) + .unwrap_or_default(); + let calls = found + .store + .tool_calls_for_session(&found.session.id) + .unwrap_or_default(); + let mut per_turn: HashMap = HashMap::new(); + for call in &calls { + let slot = per_turn.entry(call.turn_seq).or_default(); + slot.0 += 1; + slot.1 |= recorded_call_is_external(call); + } + let mut turns: Vec<(i64, TranscriptTurn)> = messages + .iter() + .filter(|m| m.mode == Mode::Text) + .filter_map(|m| { + let body = found + .store + .message_body(m) + .unwrap_or_else(|_| m.preview.clone()); + let text = atlas_agent_transcript::strip_injected_context(&body); + if text.trim().is_empty() { + return None; + } + let role = match m.role { + Role::User => "user", + Role::Assistant => "assistant", + Role::System => "system", + }; + Some(( + m.turn_seq, + TranscriptTurn { + role: role.to_string(), + text, + tool_calls: 0, + external: per_turn.get(&m.turn_seq).is_some_and(|t| t.1), + }, + )) + }) + .collect(); + // Each turn's calls, on its last assistant message. + let mut counted: std::collections::HashSet = std::collections::HashSet::new(); + for (turn_seq, turn) in turns.iter_mut().rev() { + if turn.role == "assistant" && counted.insert(*turn_seq) { + turn.tool_calls = per_turn.get(turn_seq).map_or(0, |t| t.0); + } + } + turns.into_iter().map(|(_, t)| t).collect() +} + +/// Whether a recorded call read outside content ([`is_external`]). +fn recorded_call_is_external(call: &atlas_checkpoint::ToolCall) -> bool { + call.tool_name == atlas_checkpoint::tools::ToolName::Fetch + || is_external_parts( + call.tool_name.as_str(), + call.title.as_deref(), + call.kind.as_deref(), + ) +} + /// Atlas's own tool servers: their results are not outside content. const OWN_SERVERS: [&str; 4] = ["atlas_memory", "atlas_code", "atlas_ui", "atlas_org"]; @@ -296,17 +628,22 @@ const OWN_SERVERS: [&str; 4] = ["atlas_memory", "atlas_code", "atlas_ui", "atlas /// are spelled `mcp____` by ACP agents and `.` by /// the native agent, in the tool name or the title. fn is_external(call: &atlas_agent_wire::ToolCall) -> bool { - if call.kind.as_deref() == Some("fetch") { + is_external_parts(&call.tool_name, call.title.as_deref(), call.kind.as_deref()) +} + +/// [`is_external`] over a call's name, title and ACP kind. +fn is_external_parts(tool_name: &str, title: Option<&str>, kind: Option<&str>) -> bool { + if kind == Some("fetch") { return true; } // The native agent's MCP calls carry kind `other`. A shell command, read // or edit whose title starts with a dotted word (`python3.12 -m pytest`, // `Cargo.toml`) is not one. let dotted = !matches!( - call.kind.as_deref(), + kind, Some("execute" | "read" | "edit" | "search" | "delete" | "move") ); - let mut names = std::iter::once(call.tool_name.as_str()).chain(call.title.as_deref()); + let mut names = std::iter::once(tool_name).chain(title); names.any(|name| { let lower = name.to_ascii_lowercase(); let first = lower @@ -449,6 +786,7 @@ fn completion_text(body: &str) -> Option { #[cfg(test)] mod tests { use super::*; + use crate::commands::memory_capture::Reach; use crate::commands::shared_memory::MemoryChanged; use atlas_memory::record::EntryKind; @@ -775,6 +1113,431 @@ mod tests { ); } + /// Most sessions end at quit: the end is recorded on the way out and its + /// pass is cut off with the process, so the turns kept for it are gone. + /// A short session (below the turn gates) then left nothing, and memory + /// stopped growing. The catch-up pass runs it later over the persisted + /// conversation, once. + #[tokio::test] + async fn a_session_whose_end_pass_was_lost_is_caught_up_once() { + let h = harness("catch-up", true); + // A short session: below the turn gates, its end pass never ran. + let recorded = h + .extractor + .turn_finished(&h.sharing, &h.project, &writer(), session(4)) + .await; + assert_eq!(recorded, 0); + // The app quit: the turns held for the end pass die with it. + let restarted = Extractor::new(h.memory.clone(), h.model.clone()); + assert_eq!( + restarted + .session_ended(&h.sharing, &h.project, &writer()) + .await, + 0, + "no turns held: the plain end pass has nothing to send" + ); + assert!(!restarted.holds_turns("sess-1")); + + let recorded = restarted + .catch_up(&h.sharing, &h.project, &writer(), session(4)) + .await; + assert_eq!(recorded, 4); + let decisions = h.memory.list_entries(&h.project, Some(EntryKind::Decision)); + assert_eq!(decisions.len(), 1); + assert_eq!( + decisions[0].state, + atlas_memory::record::State::Active, + "briefed, not a candidate" + ); + assert_eq!(h.model.calls().len(), 1); + + // Nothing new since: the next health pass asks nothing. + restarted + .catch_up(&h.sharing, &h.project, &writer(), session(4)) + .await; + assert_eq!(h.model.calls().len(), 1); + } + + #[tokio::test] + async fn a_session_whose_turns_are_held_here_is_not_caught_up() { + let h = harness("catch-up-live", true); + h.extractor + .turn_finished(&h.sharing, &h.project, &writer(), session(4)) + .await; + assert!(h.extractor.holds_turns("sess-1")); + let recorded = h + .extractor + .catch_up(&h.sharing, &h.project, &writer(), session(4)) + .await; + assert_eq!(recorded, 0); + assert!(h.model.calls().is_empty()); + } + + /// The persisted conversation is text only; whether the session read + /// outside content comes from the recorder, and is assumed when the + /// recorder never saw the session. + #[test] + fn the_persisted_conversation_is_read_and_unknown_tools_count_as_outside() { + let config = + std::env::temp_dir().join(format!("atlas-extract-cfg-{}", uuid::Uuid::new_v4())); + let cwd = "/nowhere/project"; + let dir = super::super::agent_transcript::dir_for(&config, cwd); + std::fs::create_dir_all(&dir).unwrap(); + std::fs::write( + dir.join("sess-9.json"), + serde_json::json!({ + "id": "sess-9", "plugin_id": "claude-acp", "cwd": cwd, + "created_at": "", "updated_at": "", + "messages": [ + {"role": "user", "content": "use bun", "timestamp": ""}, + {"role": "assistant", "content": "Done.", "timestamp": ""} + ] + }) + .to_string(), + ) + .unwrap(); + let reader = super::super::memory_capture::CaptureReader::default(); + let turns = persisted_turns(&config, cwd, "sess-9", &reader, Reach::Hosted).unwrap(); + assert_eq!( + turns + .iter() + .map(|t| (t.role.as_str(), t.text.as_str())) + .collect::>(), + [("user", "use bun"), ("assistant", "Done.")] + ); + assert!( + turns.iter().any(|t| t.external), + "no recorder: outside content assumed" + ); + // Found from another project directory too (a worktree's session). + let turns = + persisted_turns(&config, "/elsewhere", "sess-9", &reader, Reach::Hosted).unwrap(); + assert_eq!(turns.len(), 2); + assert!(persisted_turns(&config, cwd, "missing", &reader, Reach::Hosted).is_none()); + let _ = std::fs::remove_dir_all(&config); + } + + /// Record a session the way the terminal-gap importer files one: under + /// `external_jsonl`, every user entry opening a turn. + fn import_terminal_session( + project: &str, + native_id: &str, + entries: &[(atlas_checkpoint::Role, &str)], + tool: Option<(atlas_checkpoint::tools::ToolName, &str)>, + ) { + use atlas_checkpoint::model::ProjectMode; + use atlas_checkpoint::{Capture, Mode, Role, SessionKey, Source, Store, TurnContent}; + let mut store = Store::open(atlas_checkpoint::atlas_dir(project)).unwrap(); + let mut capture = Capture::new(&mut store, ProjectMode::Local); + let key = SessionKey { + workspace_id: project.to_string(), + source: Source::ExternalJsonl, + native_session_id: native_id.into(), + }; + let (mut turn, mut row) = (0i64, String::new()); + for (i, (role, body)) in entries.iter().enumerate() { + if *role == Role::User { + turn += 1; + row = capture + .record_prompt(&key, body, turn, Some("claude-code"), None, Some(project)) + .unwrap(); + } else { + capture + .record_turn( + &row, + TurnContent { + turn_seq: turn, + native_message_id: Some(format!("{native_id}-{i}")), + role: *role, + mode: Mode::Text, + body: body.to_string(), + created_at: None, + }, + ) + .unwrap(); + } + } + if let Some((name, kind)) = tool { + capture + .record_tool_call( + &row, + atlas_checkpoint::ToolCallContent { + turn_seq: turn, + native_call_id: Some("call-1"), + tool_name: name, + title: None, + kind: Some(kind), + status: atlas_checkpoint::ToolStatus::Completed, + locations: &serde_json::json!([]), + arguments: None, + result: None, + }, + ) + .unwrap(); + } + } + + /// Most work runs in terminal sessions Atlas never hosts: the recorder + /// imports their transcripts, and that is the only conversation Atlas + /// has of them. Their end pass reads it. + #[test] + fn a_terminal_sessions_conversation_is_read_from_the_recorder() { + use atlas_checkpoint::tools::ToolName; + use atlas_checkpoint::Role; + let project = super::super::memory_pack::test_support::scratch_project("terminal"); + let config = std::path::PathBuf::from( + super::super::memory_pack::test_support::scratch_project("no-transcripts"), + ); + import_terminal_session( + &project, + "cli-1", + &[ + (Role::User, "use bun, never npm"), + (Role::Assistant, "Noted: bun it is."), + (Role::User, "and run the gates"), + (Role::Assistant, "Lint, format and typecheck are green."), + ], + Some((ToolName::Bash, "execute")), + ); + let reader = super::super::memory_capture::CaptureReader::default(); + let turns = persisted_turns(&config, &project, "cli-1", &reader, Reach::WithImported) + .expect("recorded"); + assert_eq!( + turns + .iter() + .map(|t| (t.role.as_str(), t.text.as_str())) + .collect::>(), + [ + ("user", "use bun, never npm"), + ("assistant", "Noted: bun it is."), + ("user", "and run the gates"), + ("assistant", "Lint, format and typecheck are green."), + ] + ); + assert_eq!(turns.iter().map(|t| t.tool_calls).sum::(), 1); + assert!( + !turns.iter().any(|t| t.external), + "a shell call is not outside content" + ); + + // A web fetch in it makes its entries candidates. + import_terminal_session( + &project, + "cli-2", + &[ + (Role::User, "read the docs"), + (Role::Assistant, "Read them."), + ], + Some((ToolName::Fetch, "fetch")), + ); + let turns = + persisted_turns(&config, &project, "cli-2", &reader, Reach::WithImported).unwrap(); + assert!(turns.iter().any(|t| t.external)); + + // And the scan finds both, newest first, and not a live one. + let since = chrono::Utc::now().timestamp_millis() - 60_000; + let found: Vec<(String, String)> = reader + .stores(&project) + .recent_sessions(since, i64::MAX) + .into_iter() + .filter(|w| w.session_id != "live") + .map(|w| (w.agent, w.session_id)) + .collect(); + assert_eq!( + found, + [ + ("claude-code".to_string(), "cli-2".to_string()), + ("claude-code".to_string(), "cli-1".to_string()) + ] + ); + assert!( + reader + .stores(&project) + .recent_sessions(since, since) + .is_empty(), + "a session active after the idle cutoff is still running" + ); + let _ = std::fs::remove_dir_all(&project); + let _ = std::fs::remove_dir_all(&config); + } + + /// A terminal session the recorder imported from outside Atlas is read + /// for extraction, and so sent to the model, only once the user turned + /// `fromExternalSessions` on: with it off the catch-up scan never lists + /// it and its conversation is never read. + #[tokio::test] + async fn an_imported_session_is_extracted_only_when_the_user_allows_it() { + use atlas_checkpoint::Role; + let h = harness("imported-gate", true); + let config = std::path::PathBuf::from( + super::super::memory_pack::test_support::scratch_project("imported-gate-cfg"), + ); + import_terminal_session( + &h.project, + "cli-ext", + &[ + (Role::User, "we decided to use bun, never npm"), + (Role::Assistant, "Noted: bun it is."), + (Role::User, "and the ledger stays on Postgres"), + (Role::Assistant, "Kept it on Postgres."), + ], + None, + ); + let reader = super::super::memory_capture::CaptureReader::default(); + let later = chrono::Utc::now().timestamp_millis() + CATCH_UP_IDLE_MS + 60_000; + let scan = |reach| -> Vec { + recorded_catch_up_sessions(&reader.stores(&h.project), reach, later, &|_| false) + .into_iter() + .map(|w| w.session_id) + .collect() + }; + let imported = Writer { + agent: "claude-code".into(), + session_id: "cli-ext".into(), + }; + + // Off (the default): not listed, not read, nothing sent. + let reach = h.sharing.extraction_reach(&h.project); + assert_eq!(reach, Reach::Hosted); + assert!(scan(reach).is_empty(), "an imported session is not scanned"); + assert!(persisted_turns(&config, &h.project, "cli-ext", &reader, reach).is_none()); + assert!(h.model.calls().is_empty()); + + // On: listed, read and extracted. + h.sharing + .set_from_external_sessions(&h.project, true) + .unwrap(); + let reach = h.sharing.extraction_reach(&h.project); + assert_eq!(reach, Reach::WithImported); + assert_eq!(scan(reach), ["cli-ext"]); + let turns = persisted_turns(&config, &h.project, "cli-ext", &reader, reach) + .expect("the imported conversation is read"); + assert_eq!(turns.len(), 4); + let recorded = h + .extractor + .catch_up(&h.sharing, &h.project, &imported, turns) + .await; + assert!(recorded > 0); + assert_eq!(h.model.calls().len(), 1); + let _ = std::fs::remove_dir_all(&config); + } + + /// A long session caught up later is read in prompt-sized passes, not + /// one pass over its tail, at most [`CATCH_UP_PASSES`] per job; the + /// next job carries on where the last stopped. + #[tokio::test] + async fn a_long_session_is_caught_up_in_prompt_sized_passes() { + let h = harness("catch-up-long", true); + // Turn 3 fetched the web: only the pass that reads it proposes. + let long: Vec = (0..200) + .map(|i| TranscriptTurn { + role: if i % 2 == 0 { "user" } else { "assistant" }.into(), + text: format!("turn {i} {}", "x".repeat(300)), + tool_calls: 0, + external: i == 3, + }) + .collect(); + h.extractor + .catch_up(&h.sharing, &h.project, &writer(), long.clone()) + .await; + assert_eq!(h.model.calls().len(), CATCH_UP_PASSES); + let decision = h + .memory + .list_entries(&h.project, Some(EntryKind::Decision)) + .pop() + .unwrap(); + assert_eq!( + decision.state, + atlas_memory::record::State::Active, + "a later pass that read no outside content confirms what the first only proposed" + ); + h.extractor + .catch_up(&h.sharing, &h.project, &writer(), long.clone()) + .await; + let after_two = h.model.calls().len(); + assert_eq!(after_two, 2 * CATCH_UP_PASSES); + // Drained: 200 turns of ~310 chars is ~62k chars, ~11 passes. + for _ in 0..5 { + h.extractor + .catch_up(&h.sharing, &h.project, &writer(), long.clone()) + .await; + } + let total = h.model.calls().len(); + assert!((10..=13).contains(&total), "{total} passes"); + h.extractor + .catch_up(&h.sharing, &h.project, &writer(), long) + .await; + assert_eq!(h.model.calls().len(), total, "nothing new: no call"); + } + + #[test] + fn catch_up_looks_at_recent_sessions_that_are_not_live() { + let root = + std::env::temp_dir().join(format!("atlas-extract-sessions-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&root).unwrap(); + let store = atlas_memory::record::open_scope(&root).unwrap(); + let day = 24 * 3600 * 1000; + let now = 100 * day; + store + .session_started("old", "codex-acp", now - 30 * day) + .unwrap(); + store + .session_started("ended", "claude-acp", now - 3 * day) + .unwrap(); + store + .session_ended("ended", "claude-acp", now - 2 * day) + .unwrap(); + store + .session_started("open", "claude-acp", now - day) + .unwrap(); + store + .session_started("live", "claude-acp", now - day / 2) + .unwrap(); + let got: Vec = catch_up_sessions(&store, now, &|id| id == "live") + .into_iter() + .map(|w| w.session_id) + .collect(); + assert_eq!(got, ["open", "ended"]); + let _ = std::fs::remove_dir_all(&root); + } + + /// A session not live here may be live in another Atlas process on the + /// same data directory: one that started, or did anything, within the + /// idle cutoff is left for a later scan. + #[test] + fn catch_up_leaves_a_session_active_elsewhere_alone() { + let root = tempfile::TempDir::new().unwrap(); + let store = atlas_memory::record::open_scope(root.path()).unwrap(); + let minute = 60 * 1000; + let now = 1_000 * 24 * 60 * minute; + store + .session_started("just-started", "claude-acp", now - 5 * minute) + .unwrap(); + store + .session_started("busy", "codex-acp", now - 120 * minute) + .unwrap(); + store + .append_event( + atlas_memory::record::NewEvent { + agent: "codex-acp".into(), + session_id: "busy".into(), + kind: atlas_memory::record::EventKind::FileChanged, + key: "src/a.rs".into(), + payload: serde_json::json!({}), + }, + now - 2 * minute, + ) + .unwrap(); + store + .session_started("paused", "codex-acp", now - 60 * minute) + .unwrap(); + let got: Vec = catch_up_sessions(&store, now, &|_| false) + .into_iter() + .map(|w| w.session_id) + .collect(); + assert_eq!(got, ["paused"]); + } + #[test] fn outside_content_is_a_web_call_or_a_third_party_mcp_tool() { let call = diff --git a/src-tauri/src/commands/memory_indexer.rs b/src-tauri/src/commands/memory_indexer.rs index f3a045f83..5e6859179 100644 --- a/src-tauri/src/commands/memory_indexer.rs +++ b/src-tauri/src/commands/memory_indexer.rs @@ -58,6 +58,14 @@ pub enum Job { cwd: String, writer: super::shared_memory::Writer, }, + /// A session that is not live here may never have had its end pass (the + /// app quit before it finished): one over its persisted conversation, + /// a no-op when nothing is new since its last pass. Queued by the health + /// pass ([`super::memory_extract::Extractor::catch_up`]). + CatchUpExtraction { + cwd: String, + writer: super::shared_memory::Writer, + }, /// Offer the repository's high-confidence Facts to global memory /// (`atlas_memory::global`): run once when a project opens and after an /// extractor pass stores entries. @@ -119,6 +127,9 @@ pub struct MemoryRegistry { overflow: parking_lot::Mutex>, /// The last health pass per project. health: DashMap, + /// When each project's sessions were last scanned for a lost end pass + /// ([`queue_catch_up`]). + catch_up_scanned: DashMap, } impl MemoryRegistry { @@ -133,6 +144,7 @@ impl MemoryRegistry { provider: tokio::sync::Mutex::new(None), overflow: parking_lot::Mutex::new(Vec::new()), health: DashMap::new(), + catch_up_scanned: DashMap::new(), } } @@ -145,7 +157,9 @@ impl MemoryRegistry { match self.job_tx.try_send(job) { Ok(()) => {} Err(TrySendError::Full(job)) => match job { - Job::ExtractSession { .. } | Job::SessionEnded { .. } => { + Job::ExtractSession { .. } + | Job::SessionEnded { .. } + | Job::CatchUpExtraction { .. } => { tracing::warn!( target: "atlas::memory_indexer", "extraction job dropped: queue full" @@ -463,7 +477,10 @@ impl MemoryIndexer { while let Some(job) = rx.recv().await { let extraction = matches!( job, - Job::ExtractSession { .. } | Job::SessionEnded { .. } | Job::Dream { .. } + Job::ExtractSession { .. } + | Job::SessionEnded { .. } + | Job::CatchUpExtraction { .. } + | Job::Dream { .. } ); let (app, registry_for_job) = (app.clone(), registry.clone()); let work: Work = Box::pin(async move { handle(&app, ®istry_for_job, job).await }); @@ -532,6 +549,16 @@ async fn handle(app: &AppHandle, registry: &MemoryRegistry, job: Job) { if let Err(e) = index_one(registry, &cwd, &prov).await { tracing::warn!(target: "atlas::memory_indexer", "IndexCorpus {cwd} failed: {e}"); } + // The health pass scans at open; a long-running app scans again + // at most every CATCH_UP_RESCAN on the index nudges every + // finished turn and file change sends. + let due = registry + .catch_up_scanned + .get(&cwd) + .is_none_or(|at| at.elapsed() >= CATCH_UP_RESCAN); + if due { + queue_catch_up(app, registry, &cwd, chrono::Utc::now().timestamp_millis()).await; + } } Job::ExtractSession { cwd, writer, turns } => { extract(app, registry, &cwd, writer, Some(turns)).await; @@ -539,6 +566,9 @@ async fn handle(app: &AppHandle, registry: &MemoryRegistry, job: Job) { Job::SessionEnded { cwd, writer } => { extract(app, registry, &cwd, writer, None).await; } + Job::CatchUpExtraction { cwd, writer } => { + catch_up(app, registry, &cwd, writer).await; + } Job::Compact { cwd } => { if let Err(e) = compact_one(registry, &cwd).await { tracing::warn!(target: "atlas::memory_indexer", "Compact {cwd} failed: {e}"); @@ -685,6 +715,8 @@ async fn health_one(app: &AppHandle, registry: &MemoryRegistry, cwd: &str) -> Re recent.into_iter().map(|s| (cwd.to_string(), s)).collect(), ) .await; + // Recent sessions that are not live here: an end pass that never ran. + queue_catch_up(app, registry, cwd, now).await; // The Agent Memory Repo mirror, when the user turned it on: rewritten // and committed only when memory changed. let mirror = app @@ -851,8 +883,109 @@ async fn extract( let sharing = app.state::(); let stored = match turns { Some(turns) => extractor.turn_finished(&sharing, cwd, &writer, turns).await, - None => extractor.session_ended(&sharing, cwd, &writer).await, + None => { + let held = extractor.holds_turns(&writer.session_id); + let stored = extractor.session_ended(&sharing, cwd, &writer).await; + if held { + stored + } else { + // No turn of it finished in this process (it was resumed and + // closed, or its turns ran before a restart): the end pass + // reads what Atlas persisted instead. + return catch_up(app, registry, cwd, writer).await; + } + } + }; + reindex_after(registry, cwd, stored); +} + +/// How often a project's sessions are scanned for a lost end pass while the +/// app runs (besides the scan at open). +const CATCH_UP_RESCAN: Duration = Duration::from_secs(30 * 60); + +/// Queue an end pass for every recent session of `cwd` that is not live in +/// this process and may never have had one ([`Job::CatchUpExtraction`]; a +/// no-op for one with nothing new since its last pass). +/// +/// Two kinds: sessions Atlas hosted whose end pass died with the process +/// (most sessions end at quit), and sessions Atlas never hosted. The +/// latter are the terminal ones the capture recorder imports off disk, and +/// they are most of the work: without this, nothing they decided or learned +/// reached memory, which grew only from in-app sessions. Reading them sends +/// their conversation to the extraction model, so they are scanned only when +/// the user turned that on (`MemorySharingState::extraction_reach`). +async fn queue_catch_up(app: &AppHandle, registry: &MemoryRegistry, cwd: &str, now: i64) { + registry + .catch_up_scanned + .insert(cwd.to_string(), std::time::Instant::now()); + let Some(extractor) = app + .try_state::>() + .map(|e| e.inner().clone()) + else { + return; + }; + let reach = app + .state::() + .extraction_reach(cwd); + let (owned, reader) = (cwd.to_string(), capture_reader(app)); + let lost = tokio::task::spawn_blocking(move || { + let live = |id: &str| extractor.holds_turns(id); + let recorded = super::memory_extract::recorded_catch_up_sessions( + &reader.stores(&owned), + reach, + now, + &live, + ); + let hosted = super::shared_memory::store_for(&owned) + .map(|store| super::memory_extract::catch_up_sessions(&store, now, &live)) + .unwrap_or_default(); + super::memory_extract::merge_catch_up(recorded, hosted) + }) + .await + .unwrap_or_default(); + for writer in lost { + registry.send_or_remember(Job::CatchUpExtraction { + cwd: cwd.to_string(), + writer, + }); + } +} + +/// The end pass of a session whose turns this process does not hold, over +/// its persisted conversation (`memory_extract::persisted_turns`). +async fn catch_up( + app: &AppHandle, + registry: &MemoryRegistry, + cwd: &str, + writer: super::shared_memory::Writer, +) { + let Some(extractor) = app.try_state::>() else { + return; + }; + let Some(config_dir) = app + .try_state::>() + .map(|t| t.config_dir().to_path_buf()) + else { + return; + }; + let reader = capture_reader(app); + // Decided when the pass runs, not when it was queued: a session imported + // from outside Atlas is not read once the user turned that off. + let reach = app + .state::() + .extraction_reach(cwd); + let (owned, session) = (cwd.to_string(), writer.session_id.clone()); + let turns = tokio::task::spawn_blocking(move || { + super::memory_extract::persisted_turns(&config_dir, &owned, &session, &reader, reach) + }) + .await + .ok() + .flatten(); + let Some(turns) = turns else { + return; }; + let sharing = app.state::(); + let stored = extractor.catch_up(&sharing, cwd, &writer, turns).await; reindex_after(registry, cwd, stored); } diff --git a/src-tauri/src/commands/memory_server/briefing.rs b/src-tauri/src/commands/memory_server/briefing.rs index 958beeb8c..2fc62b2e6 100644 --- a/src-tauri/src/commands/memory_server/briefing.rs +++ b/src-tauri/src/commands/memory_server/briefing.rs @@ -228,13 +228,17 @@ pub(super) fn read_briefing( now: i64, check: &dyn Fn(&[Entry]) -> Checked, ) -> anyhow::Result { - let plan = store.list(EntryKind::Plan, 1, Origin::Any)?.pop(); - let mut files_changed = store.list( - EntryKind::FileChanged, - EntryKind::FileChanged.cap(), - Origin::Any, - )?; + let plan = store + .list(EntryKind::Plan, 1, Origin::Any)? + .pop() + .filter(is_active_plan); + // Only this scope's files, newest first: rows recorded before capture + // kept to the scope can name another checkout, `/tmp` or `~/.claude`. + let dirs = store.scope_dirs(); + let mut files_changed = store.list(EntryKind::FileChanged, RANK_POOL, Origin::Any)?; files_changed.reverse(); + files_changed.retain(|e| atlas_memory::record::path_in_dirs(&e.key, &dirs)); + files_changed.truncate(EntryKind::FileChanged.cap()); let mut durable = Vec::new(); for kind in DURABLE_KINDS { durable.extend(store.list(kind, RANK_POOL, Origin::Any)?); @@ -267,6 +271,84 @@ pub(super) fn read_briefing( }) } +/// Whether a stored plan is still the active one: its status says so, and +/// not every item of it is done. Capture keeps only the newest plan, so a +/// later session's plan has already replaced an older one; what is left to +/// catch is a plan that ran to completion and was never cleared (capture +/// used to log every plan as `active`, finished or not). +pub(super) fn is_active_plan(plan: &Entry) -> bool { + (plan.status.is_empty() || plan.status == "active") + && !atlas_memory::handoff::plan_is_finished(&plan.content) +} + +/// The handoff a briefing for session `own` carries: what the previous +/// session of any agent left, with its recorded facts read now (its commits +/// may have landed after it ended) and only this scope's files. +/// +/// Memory's own note ([`RecordStore::last_episode`]) covers the sessions +/// Atlas hosted. The capture recorder also holds the terminal sessions its +/// importer read off disk, which memory never saw start or end; when one of +/// those did work after memory's newest note, it is the previous session, +/// and its note is built from the recorder (and whatever memory logged for +/// it). Blocking (SQLite, git). +pub(super) fn read_handoff( + store: &RecordStore, + own: &str, + stores: &crate::commands::memory_capture::ScopeStores, +) -> Option { + /// How many newer recorded sessions are tried, newest first. + const TRIED: usize = 5; + let mut note = store.last_episode(own).ok().flatten(); + let since = note + .as_ref() + .map_or(i64::MIN, |n| n.ended_at.saturating_add(1)); + let known = note.as_ref().map(|n| n.session.clone()).unwrap_or_default(); + let newer = stores + .recent_sessions(since, i64::MAX) + .into_iter() + .filter(|w| w.session_id != own && w.session_id != known) + .take(TRIED); + for writer in newer { + let Some(found) = stores.find(&writer.session_id) else { + continue; + }; + let session = &found.session; + let ended_at = session + .last_activity_at + .unwrap_or(session.updated_at) + .timestamp_millis(); + let Ok(mut built) = atlas_memory::handoff::build_handoff( + store, + &writer.session_id, + &writer.agent, + ended_at, + ) else { + continue; + }; + built.started_at = built + .started_at + .or(Some(session.started_at.timestamp_millis())); + let facts = crate::commands::memory_capture::session_facts(&found); + let worth = !built.is_empty() || !facts.is_empty(); + built.apply_facts(facts); + if worth { + note = Some(built); + break; + } + } + let mut note = note?; + if note.session != own { + if let Some(found) = stores.find(¬e.session) { + note.apply_facts(crate::commands::memory_capture::session_facts(&found)); + } + } + let dirs = store.scope_dirs(); + note.files.retain(|f| { + atlas_memory::record::path_in_dirs(f.strip_suffix(" (deleted)").unwrap_or(f), &dirs) + }); + Some(note) +} + /// The preferences a briefing leads with, from the stored ones (oldest /// first): the active ones, newest first, at most [`PREFERENCES_MAX`] and /// [`PREFERENCES_MAX_CHARS`] of content, the oldest dropped first. @@ -729,6 +811,137 @@ mod briefing_tests { .is_none()); } + fn log( + store: &RecordStore, + kind: atlas_memory::record::EventKind, + key: &str, + payload: Value, + ts: i64, + ) { + store + .append_event( + atlas_memory::record::NewEvent { + agent: "claude-acp".into(), + session_id: "s-old".into(), + kind, + key: key.into(), + payload, + }, + ts, + ) + .unwrap(); + } + + fn scope(label: &str) -> (std::path::PathBuf, std::sync::Arc) { + let root = + std::env::temp_dir().join(format!("atlas-brief-{label}-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&root).unwrap(); + let store = open_scope(&root).unwrap(); + (root, store) + } + + /// A plan whose every item is done is not the active plan, even when it + /// was stored as `active` (the live store's August plan was). + #[test] + fn a_finished_plan_is_not_briefed_as_active() { + use atlas_memory::record::EventKind; + let (root, store) = scope("plan"); + log( + &store, + EventKind::PlanSet, + "plan", + json!({"text": "- [completed] Create branch for issue #177\n- [completed] Commit, push, open PR", "status": "active"}), + 1, + ); + let b = read_briefing(&store, 2, &|_| Checked::default()).unwrap(); + assert!(b.plan.is_none(), "{:?}", b.plan); + + log( + &store, + EventKind::PlanSet, + "plan", + json!({"text": "- [completed] Read auth\n- [in_progress] Move to EdDSA", "status": "active"}), + 3, + ); + let b = read_briefing(&store, 4, &|_| Checked::default()).unwrap(); + assert!(b.plan.unwrap().content.contains("Move to EdDSA")); + let _ = std::fs::remove_dir_all(&root); + } + + /// Files changed in another checkout, `/tmp` or `~/.claude` are not this + /// project's, however they got into its store. + #[test] + fn the_briefing_lists_only_files_under_the_project_root() { + use atlas_memory::record::EventKind; + let (root, store) = scope("files"); + let inside = root + .canonicalize() + .unwrap() + .join("src/main.rs") + .to_string_lossy() + .into_owned(); + for (i, path) in [ + inside.as_str(), + "/Users/someone/Developer/atlas/src/App.tsx", + "/tmp/homebrew-cask-pr/Casks/a/atlas-ai.rb", + "/Users/someone/.claude/plans/plan.md", + ] + .into_iter() + .enumerate() + { + log( + &store, + EventKind::FileChanged, + path, + json!({"path": path, "summary": "Edit"}), + i as i64 + 1, + ); + } + let b = read_briefing(&store, 10, &|_| Checked::default()).unwrap(); + let paths: Vec<&str> = b.files_changed.iter().map(|e| e.key.as_str()).collect(); + assert_eq!(paths, [inside.as_str()]); + let _ = std::fs::remove_dir_all(&root); + } + + /// The previous session was a terminal one Atlas never hosted: memory + /// has no note of it, the recorder has its work. The briefing hands + /// that on; a note memory stored later wins. + #[test] + fn a_terminal_session_the_recorder_saw_is_handed_off() { + use crate::commands::memory_capture::test_support::Recording; + use crate::commands::memory_capture::CaptureReader; + let (root, store) = scope("handoff-recorded"); + let project = root.to_string_lossy().into_owned(); + assert!(read_handoff(&store, "own", &CaptureReader::default().stores(&project)).is_none()); + + let mut rec = Recording::open_turn(&project, "cli-1", "claude-code", "move auth to EdDSA"); + rec.write("src/auth.rs", b"pub fn sign() {}\n"); + rec.close_turn(); + drop(rec); + let stores = CaptureReader::default().stores(&project); + let note = read_handoff(&store, "own", &stores).expect("a handoff"); + assert_eq!(note.session, "cli-1"); + assert_eq!(note.agent, "claude-code"); + assert_eq!(note.files, ["src/auth.rs"]); + assert_eq!(note.title.as_deref(), Some("move auth to EdDSA")); + // A session is not handed its own work. + assert!(read_handoff(&store, "cli-1", &stores).is_none()); + + // A note memory stored after it is the previous session. + store + .record_episode(&atlas_memory::handoff::HandoffNote { + session: "s-later".into(), + agent: "codex-acp".into(), + ended_at: chrono::Utc::now().timestamp_millis() + 60_000, + decisions: vec!["Keep WAL mode".into()], + ..Default::default() + }) + .unwrap(); + let note = read_handoff(&store, "own", &stores).unwrap(); + assert_eq!(note.session, "s-later"); + let _ = std::fs::remove_dir_all(&root); + } + /// More than a page's worth sharing one timestamp is returned whole, /// rather than looping forever on an unmovable clock. #[test] diff --git a/src-tauri/src/commands/memory_server/mod.rs b/src-tauri/src/commands/memory_server/mod.rs index c2f11321f..c610f44f3 100644 --- a/src-tauri/src/commands/memory_server/mod.rs +++ b/src-tauri/src/commands/memory_server/mod.rs @@ -8,10 +8,11 @@ //! - **One bearer token per (session, scope)** ([`tokens`]): minted when a //! session starts, revoked when it ends, checked on every request. The //! token says who is calling and which scope's record to open. -//! - **Handed to every agent** ([`offers`]): ACP agents that advertise -//! `mcpCapabilities.http` over HTTP, the others as a stdio server (this -//! binary's `mcp-bridge`, the token in its environment; ADR-0019), and the -//! native agent through its thread's engine config. +//! - **Handed to every agent** ([`offers`]): every ACP agent as a stdio +//! server (this binary's `mcp-bridge`, ADR-0019), the token in a private +//! file the entry names and never in the entry itself, which an adapter may +//! put on a command line (ADR-0020); the native agent over HTTP through its +//! thread's engine config. //! - **Ten tools** ([`tools`]), read first, write last: `memory_briefing`, //! `memory_changes`, `memory_search`, `memory_get`, `memory_list`, //! `memory_history`, `memory_why`, `memory_remember`, `memory_feedback`, @@ -29,7 +30,7 @@ mod bench; mod briefing; mod host; -mod offers; +pub(crate) mod offers; #[cfg(test)] mod tests; mod tokens; @@ -42,7 +43,7 @@ pub use briefing::{SessionClocks, SessionReads}; #[allow(unused_imports)] pub use host::{MemoryServer, MemoryServerHost, SharingGate, Sources}; #[allow(unused_imports)] -pub use offers::{MemorySessionOffers, OfferDecision, BRIDGE_ARG, BRIDGE_TOKEN_ENV}; +pub use offers::{MemorySessionOffers, OfferDecision, BRIDGE_ARG}; #[allow(unused_imports)] pub use tokens::{Grant, MemoryTokens}; /// How long an agent may cache a `tools/list` answer; the UI tool server diff --git a/src-tauri/src/commands/memory_server/offers.rs b/src-tauri/src/commands/memory_server/offers.rs index 404b834ae..83c6bd0df 100644 --- a/src-tauri/src/commands/memory_server/offers.rs +++ b/src-tauri/src/commands/memory_server/offers.rs @@ -1,14 +1,14 @@ //! Handing the server to sessions. //! //! Every agent that can take the server is handed it on each session request -//! ([`MemorySessionOffers`]): an ACP session request carries the server in -//! `mcpServers` when the agent advertised `mcpCapabilities.http`; an agent -//! that did not gets it as a stdio server, the Atlas binary's own -//! `mcp-bridge` forwarding to the loopback server with the token in its -//! environment (ADR-0019); the native agent gets it as a StreamableHttp -//! entry in its thread's engine config. The -//! token is minted for the *request*, before a new session's id exists, and -//! bound to the id once the agent answers; an offer that never binds is +//! ([`MemorySessionOffers`]). An agent in another process — every ACP agent, +//! whatever it advertised — gets it as a stdio server, the Atlas binary's own +//! `mcp-bridge` forwarding to the loopback server (ADR-0019), with the token +//! in a private file the entry names, never in the entry itself: an adapter +//! may put the entry on a command line (ADR-0020). The native agent, in this +//! process, gets it as a StreamableHttp entry with the token in its header. +//! The token is minted for the *request*, before a new session's id exists, +//! and bound to the id once the agent answers; an offer that never binds is //! revoked. Each decision is logged, one line per session request. use std::sync::Arc; @@ -19,6 +19,7 @@ use atlas_agent_servers::{AskFirst, SessionMcpOffer, SessionMcpRequest, SessionM use super::host::{MemoryServerHost, SharingGate}; use super::MEMORY_SERVER_NAME; use crate::commands::code_server::{CodeOffer, CodeOfferDecision, CODE_PATH, CODE_SERVER_NAME}; +use crate::commands::memory_bridge::TokenFiles; use crate::commands::org_server::{ OrgOffer, OrgOfferDecision, EVERY_TIME_TOOLS, ORG_PATH, ORG_SERVER_NAME, OUTWARD_TOOLS, }; @@ -27,9 +28,10 @@ use crate::commands::ui_server::{UiOffer, UiOfferDecision, UI_PATH, UI_SERVER_NA /// Whether one session request is handed the memory tool server. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum OfferDecision { + /// Included over HTTP: the agent runs in this process. Included, - /// Included as a stdio server through `atlas mcp-bridge`: the agent did - /// not advertise HTTP MCP (ADR-0019). + /// Included as a stdio server through `atlas mcp-bridge`: the agent runs + /// in another process (ADR-0019, ADR-0020). IncludedViaBridge, /// Left out, and why. Omitted(&'static str), @@ -51,15 +53,16 @@ impl OfferDecision { } /// Included in a project with shared memory on (the tools would hold - /// nothing otherwise), once the server is running: over HTTP for an agent - /// that advertised HTTP MCP, else through the stdio bridge. Never decided - /// by which agent it is. - pub fn decide(http_mcp: bool, sharing_on: bool, server_running: bool) -> Self { + /// nothing otherwise), once the server is running: over HTTP when the + /// entry may carry the token (`over_http`: an in-process agent that takes + /// HTTP), else through the stdio bridge. Never decided by which agent it + /// is. + pub fn decide(over_http: bool, sharing_on: bool, server_running: bool) -> Self { if !sharing_on { Self::Omitted("shared memory is off for this project") } else if !server_running { Self::Omitted("memory tool server is not running") - } else if http_mcp { + } else if over_http { Self::Included } else { Self::IncludedViaBridge @@ -69,18 +72,21 @@ impl OfferDecision { /// The argument that makes the Atlas binary the stdio bridge. pub const BRIDGE_ARG: &str = "mcp-bridge"; -/// The environment variable the bridge reads its session token from (never -/// argv, which other users' `ps` can read). -pub const BRIDGE_TOKEN_ENV: &str = "ATLAS_MCP_TOKEN"; -/// `name` at `url` as a stdio server: ` mcp-bridge `, the -/// token in the environment. -fn bridge_entry(exe: &std::path::Path, name: &str, url: String, token: &str) -> acp::McpServer { - acp::McpServer::Stdio( - acp::McpServerStdio::new(name, exe.to_path_buf()) - .args(vec![BRIDGE_ARG.to_string(), url]) - .env(vec![acp::EnvVariable::new(BRIDGE_TOKEN_ENV, token)]), - ) +/// `name` at `url` as a stdio server: ` mcp-bridge `. Nothing in the entry is secret: an adapter may put all of it, +/// `env` included, on a command line (ADR-0020). +fn bridge_entry( + exe: &std::path::Path, + name: &str, + url: String, + token_file: &std::path::Path, +) -> acp::McpServer { + acp::McpServer::Stdio(acp::McpServerStdio::new(name, exe.to_path_buf()).args(vec![ + BRIDGE_ARG.to_string(), + url, + token_file.to_string_lossy().into_owned(), + ])) } impl OfferDecision { @@ -104,9 +110,12 @@ pub struct MemorySessionOffers { ui: Option, org: Option, code: Option, + token_files: Arc, } impl MemorySessionOffers { + /// Bridged sessions' token files go to a fresh temporary directory until + /// [`with_token_files`](Self::with_token_files) names the app's own. pub fn new(host: Arc, gate: SharingGate) -> Self { Self { host, @@ -114,9 +123,17 @@ impl MemorySessionOffers { ui: None, org: None, code: None, + token_files: Arc::new(TokenFiles::in_temp()), } } + /// Write bridged sessions' token files with `files` (ADR-0020), shared + /// with whatever removes a session's file when the session ends. + pub fn with_token_files(mut self, files: Arc) -> Self { + self.token_files = files; + self + } + /// Also offer the UI tool server, mounted on this host at `/ui`. pub fn with_ui(mut self, ui: UiOffer) -> Self { self.ui = Some(ui); @@ -179,21 +196,25 @@ impl SessionMcpServers for MemorySessionOffers { fn offer(&self, request: &SessionMcpRequest) -> SessionMcpOffer { let cwd = request.cwd.to_string_lossy().into_owned(); let agent = request.agent_id.as_str().to_string(); - // An agent without HTTP MCP reaches the servers through this binary - // as a stdio bridge (ADR-0019); without a path to it, it gets none. - let bridge = if request.http_mcp { + // Only an agent in this process may hold the token in its entry: any + // other may put the entry on a command line (ADR-0020), so it reaches + // the servers through this binary as a stdio bridge (ADR-0019) that + // reads the token from a private file. Without a path to the binary, + // it gets none. + let over_http = request.in_process && request.http_mcp; + let bridge = if over_http { None } else { std::env::current_exe().ok() }; - let reachable = request.http_mcp || bridge.is_some(); + let reachable = over_http || bridge.is_some(); // The gate reads the sharing file; only asked when it can matter. let sharing_on = reachable && (self.gate)(&cwd); let url = self.host.url(); let decision = if reachable { - OfferDecision::decide(request.http_mcp, sharing_on, url.is_some()) + OfferDecision::decide(over_http, sharing_on, url.is_some()) } else { - OfferDecision::Omitted("agent did not advertise mcpCapabilities.http") + OfferDecision::Omitted("no transport reaches the agent") }; tracing::info!( target: "atlas::memory_server", @@ -271,11 +292,29 @@ impl SessionMcpServers for MemorySessionOffers { // org decision names none otherwise). let tokens = self.host.tokens().clone(); let token = tokens.mint_unbound(&agent, &cwd, scope, ui_included); + // A bridged session's entries name one file holding the token. + let token_file = match &bridge { + Some(_) => match self + .token_files + .write(&token, |t| tokens.grant(t).is_some()) + { + Ok(path) => Some(path), + Err(e) => { + tracing::warn!( + target: "atlas::memory_server", + "no Atlas tool servers for this session: cannot write its token file: {e:#}" + ); + tokens.revoke_token(&token); + return SessionMcpOffer::none(); + } + }, + None => None, + }; let servers = entries .into_iter() - .map(|(name, url)| match &bridge { - Some(exe) => bridge_entry(exe, name, url, &token), - None => acp::McpServer::Http(acp::McpServerHttp::new(name, url).headers(vec![ + .map(|(name, url)| match (&bridge, &token_file) { + (Some(exe), Some(file)) => bridge_entry(exe, name, url, file), + _ => acp::McpServer::Http(acp::McpServerHttp::new(name, url).headers(vec![ acp::HttpHeader::new("Authorization", format!("Bearer {token}")), ])), }) @@ -290,10 +329,55 @@ impl SessionMcpServers for MemorySessionOffers { } else { AskFirst::none() }; + let token_files = self.token_files.clone(); SessionMcpOffer::new(servers, move |session| match session { - Some(id) => tokens.bind(&token, &id.to_string()), - None => tokens.revoke_token(&token), + Some(id) => { + let id = id.to_string(); + tokens.bind(&token, &id); + // Its file goes when the session ends. + if let Some(file) = &token_file { + token_files.bind(file, &id); + } + } + None => { + tokens.revoke_token(&token); + if let Some(file) = &token_file { + token_files.remove(file); + } + } }) .asking_first(ask_first) } } + +/// Every entry `offer` carries, as `(name, url, bearer token)`: the token +/// from an HTTP entry's header, or from the file a bridged entry names. +#[cfg(test)] +pub fn offered_entries(offer: &SessionMcpOffer) -> Vec<(String, String, String)> { + offer + .servers() + .iter() + .map(|server| match server { + acp::McpServer::Http(http) => { + let token = http + .headers + .iter() + .find(|h| h.name == "Authorization") + .and_then(|h| h.value.strip_prefix("Bearer ")) + .expect("a bearer token") + .to_string(); + (http.name.clone(), http.url.clone(), token) + } + acp::McpServer::Stdio(stdio) => { + let [arg, url, file] = stdio.args.as_slice() else { + panic!("mcp-bridge : {:?}", stdio.args) + }; + assert_eq!(arg, BRIDGE_ARG); + let token = crate::commands::memory_bridge::read_token_file(file.as_ref()) + .expect("a readable token file"); + (stdio.name.clone(), url.clone(), token) + } + other => panic!("{other:?}"), + }) + .collect() +} diff --git a/src-tauri/src/commands/memory_server/tests.rs b/src-tauri/src/commands/memory_server/tests.rs index b39a6c43f..a9f97af92 100644 --- a/src-tauri/src/commands/memory_server/tests.rs +++ b/src-tauri/src/commands/memory_server/tests.rs @@ -1176,33 +1176,131 @@ async fn running_host(gate: SharingGate) -> Arc { host } +/// An ACP agent's request: another process, whatever it advertised. fn request(http_mcp: bool, cwd: &str, session: Option<&str>) -> SessionMcpRequest { SessionMcpRequest { agent_id: atlas_acp_thread::AgentId::new("claude-code"), http_mcp, ui_control: false, org_access: false, + in_process: false, cwd: std::path::PathBuf::from(cwd), session_id: session.map(acp::SessionId::new), } } -/// The one server an offer carries, as `(name, url, bearer token)`. +/// The native agent's request: in this process, over HTTP. +fn in_process_request(cwd: &str) -> SessionMcpRequest { + SessionMcpRequest { + agent_id: atlas_acp_thread::AgentId::new("atlas-agent"), + in_process: true, + ..request(true, cwd, None) + } +} + +/// The one server an offer carries, as `(name, url, bearer token)`: from the +/// header of an HTTP entry, or from the file a bridged entry names. fn offered(offer: &SessionMcpOffer) -> Option<(String, String, String)> { - match offer.servers() { - [acp::McpServer::Http(http)] => { - let token = http - .headers - .iter() - .find(|h| h.name == "Authorization") - .and_then(|h| h.value.strip_prefix("Bearer ")) - .expect("the entry carries a bearer token") - .to_string(); - Some((http.name.clone(), http.url.clone(), token)) + let mut entries = offers::offered_entries(offer); + assert!(entries.len() <= 1, "one server at most: {entries:?}"); + entries.pop() +} + +/// ADR-0020: an agent in another process may put its whole MCP entry on a +/// command line (the Claude Agent SDK's `--mcp-config `), so no field +/// of any entry it is offered carries the token — not a header, not the +/// URL, not an argument, not `env` — whether or not it advertised HTTP MCP. +/// The token is in the private file the entry names. Only the in-process +/// agent's entry holds the token itself, in its header. +#[tokio::test(flavor = "multi_thread")] +async fn no_entry_offered_to_another_process_carries_the_token() { + let host = running_host(always_on()).await; + let offers = MemorySessionOffers::new(host.clone(), always_on()).with_code( + crate::commands::code_server::CodeOffer::new(Arc::new(|| true)), + ); + for http_mcp in [true, false] { + let offer = offers.offer(&request(http_mcp, "/p", None)); + assert_eq!( + offer.servers().len(), + 2, + "memory and code: {:?}", + offer.servers() + ); + let mut tokens = Vec::new(); + for server in offer.servers() { + let acp::McpServer::Stdio(stdio) = server else { + panic!("http_mcp={http_mcp}: only the stdio bridge leaves the process: {server:?}") + }; + assert!(stdio.env.is_empty(), "nothing in env: {:?}", stdio.env); + let file = stdio.args.last().unwrap(); + let token = crate::commands::memory_bridge::read_token_file(file.as_ref()).unwrap(); + assert!(host.tokens().grant(&token).is_some(), "a live token"); + let printed = serde_json::to_string(server).unwrap(); + assert!( + !printed.contains(&token), + "http_mcp={http_mcp}: the token is in the entry: {printed}" + ); + tokens.push(token); } - [] => None, - other => panic!("one server at most: {other:?}"), + assert_eq!(tokens[0], tokens[1], "one token for both servers"); } + + let offer = offers.offer(&in_process_request("/p")); + let token = host.tokens().grant(&offered_token(&offer)); + assert!( + token.is_some(), + "the in-process agent's header token is live" + ); +} + +/// The token of an offer whose entries are all HTTP. +fn offered_token(offer: &SessionMcpOffer) -> String { + let [acp::McpServer::Http(http), ..] = offer.servers() else { + panic!("HTTP for the in-process agent: {:?}", offer.servers()) + }; + http.headers[0] + .value + .strip_prefix("Bearer ") + .unwrap() + .to_string() +} + +/// An offer released unbound removes its token file with its token. +#[tokio::test(flavor = "multi_thread")] +async fn a_bridged_offer_that_never_binds_removes_its_token_file() { + let host = running_host(always_on()).await; + let offers = MemorySessionOffers::new(host, always_on()); + let offer = offers.offer(&request(true, "/p", None)); + let [acp::McpServer::Stdio(stdio)] = offer.servers() else { + panic!("{:?}", offer.servers()) + }; + let file = std::path::PathBuf::from(stdio.args.last().unwrap()); + assert!(file.exists()); + drop(offer); + assert!(!file.exists(), "the file goes with the token"); +} + +/// A bound offer's token file goes when its session ends, not at some later +/// offer's sweep. +#[tokio::test(flavor = "multi_thread")] +async fn a_bound_sessions_token_file_goes_when_the_session_ends() { + let host = running_host(always_on()).await; + let files = Arc::new(crate::commands::memory_bridge::TokenFiles::in_temp()); + let offers = MemorySessionOffers::new(host, always_on()).with_token_files(files.clone()); + let offer = offers.offer(&request(true, "/p", None)); + let [acp::McpServer::Stdio(stdio)] = offer.servers() else { + panic!("{:?}", offer.servers()) + }; + let file = std::path::PathBuf::from(stdio.args.last().unwrap()); + offer.bind(&acp::SessionId::new("s-bound")); + assert!( + file.exists(), + "a bound session keeps its file while it runs" + ); + files.session_ended("s-other"); + assert!(file.exists()); + files.session_ended("s-bound"); + assert!(!file.exists(), "the file goes with the session"); } #[tokio::test(flavor = "multi_thread")] @@ -1254,33 +1352,33 @@ async fn an_http_agent_is_offered_the_server_with_a_token_that_binds_to_its_sess let _ = std::fs::remove_dir_all(&project); } -/// ADR-0019: an agent without HTTP MCP is handed the server as a stdio -/// entry, this binary's `mcp-bridge`, with the token in its environment and -/// never on its command line. +/// ADR-0019, ADR-0020: an ACP agent is handed the server as a stdio entry, +/// this binary's `mcp-bridge`, naming a file that holds the token. #[tokio::test(flavor = "multi_thread")] -async fn an_agent_without_http_mcp_is_offered_the_stdio_bridge() { +async fn an_acp_agent_is_offered_the_stdio_bridge() { let host = running_host(always_on()).await; let offers = MemorySessionOffers::new(host.clone(), always_on()); let offer = offers.offer(&request(false, "/p", None)); let [acp::McpServer::Stdio(stdio)] = offer.servers() else { panic!("{:?}", offer.servers()) }; - assert_eq!(stdio.name, MEMORY_SERVER_NAME); - assert_eq!(stdio.args, [BRIDGE_ARG.to_string(), host.url().unwrap()]); - let token = stdio - .env - .iter() - .find(|v| v.name == BRIDGE_TOKEN_ENV) - .map(|v| v.value.clone()) - .expect("the token rides the environment"); - assert!(!token.is_empty()); - assert!( - stdio.args.iter().all(|a| !a.contains(&token)), - "no token in argv" - ); + assert_eq!(stdio.command, std::env::current_exe().unwrap()); + let (name, url, token) = offered(&offer).unwrap(); + assert_eq!(name, MEMORY_SERVER_NAME); + assert_eq!(Some(url), host.url()); assert!(host.tokens().grant(&token).is_some(), "a live token"); } +/// The in-process agent keeps HTTP, its token in the header it reads from +/// memory. +#[tokio::test(flavor = "multi_thread")] +async fn the_in_process_agent_is_offered_the_server_over_http() { + let host = running_host(always_on()).await; + let offers = MemorySessionOffers::new(host, always_on()); + let offer = offers.offer(&in_process_request("/p")); + assert!(matches!(offer.servers(), [acp::McpServer::Http(_)])); +} + /// A stdio-only agent's messages reach the loopback server through the /// bridge, with the MCP session id carried from the first response on. #[tokio::test(flavor = "multi_thread")] @@ -1377,6 +1475,11 @@ fn each_decision_is_one_log_line_naming_the_agent_its_capability_and_the_outcome OfferDecision::decide(false, true, true).log_line("gemini", false), "memory tool server offer: agent=gemini http_mcp=false memory_server=included_via_bridge", ); + // An ACP agent that advertised HTTP is bridged all the same (ADR-0020). + assert_eq!( + OfferDecision::decide(false, true, true).log_line("claude-code", true), + "memory tool server offer: agent=claude-code http_mcp=true memory_server=included_via_bridge", + ); assert_eq!( OfferDecision::decide(true, false, true).log_line("atlas-agent", true), "memory tool server offer: agent=atlas-agent http_mcp=true memory_server=omitted \ @@ -1698,6 +1801,34 @@ async fn the_handoff_carries_what_the_recorded_session_did() { let _ = std::fs::remove_dir_all(&p); } +/// The description promises a `handoff` field, so it is there even when no +/// earlier session left a note: `null`, never missing. +#[tokio::test(flavor = "multi_thread")] +async fn a_briefing_with_no_earlier_session_carries_a_null_handoff() { + let p = temp_project("handoff-none"); + let memory = ticking_memory(); + let tokens = Arc::new(MemoryTokens::default()); + let server = serve( + memory.clone(), + tokens.clone(), + always_on(), + Sources::default(), + ) + .await; + let b = connect(&server.url(), &tokens.mint("s-only", "codex", &p)) + .await + .unwrap(); + let (err, brief) = call(&b, "memory_briefing", json!({})).await; + assert!(!err, "{brief}"); + assert!( + brief.as_object().unwrap().contains_key("handoff"), + "{brief}" + ); + assert!(brief["handoff"].is_null(), "{brief}"); + b.cancel().await.ok(); + let _ = std::fs::remove_dir_all(&p); +} + #[test] fn a_long_handoff_is_cut_files_first_and_decisions_last() { let note = atlas_memory::handoff::HandoffNote { @@ -1888,3 +2019,169 @@ async fn memory_why_marks_a_memory_whose_cited_lines_changed_as_stale() { a.cancel().await.ok(); let _ = std::fs::remove_dir_all(&p); } + +/// A server whose index answers every query with `docs`. +async fn serve_with_documents(docs: Vec) -> (MemoryServer, Arc) { + let tokens = Arc::new(MemoryTokens::default()); + let index: IndexSearch = Arc::new(move |_cwd, _query, limit: usize| { + let docs: Vec = docs.iter().take(limit).cloned().collect(); + Box::pin(async move { docs }) + }); + let server = serve( + ticking_memory(), + tokens.clone(), + always_on(), + Sources { + index: Some(index), + ..Sources::default() + }, + ) + .await; + (server, tokens) +} + +/// The live miss: one query returned ~4k tokens, two of its documents whole +/// session transcripts. A document is an excerpt around the passage that +/// matched, not the transcript's head; the reply stays inside its budget, +/// and says so when it had to drop results. +#[tokio::test(flavor = "multi_thread")] +async fn memory_search_documents_are_excerpts_inside_the_output_budget() { + let project = temp_project("doc-budget"); + let filler = "the agent read some files and ran the tests again. ".repeat(600); + let transcript = + format!("SESSION HEAD {filler}we built the grep prefilter for large repositories {filler}"); + let docs: Vec = (0..6) + .map(|i| IndexDoc { + id: Some(format!("claude:session-{i}")), + title: format!("session {i}"), + source: "claude".to_string(), + text: transcript.clone(), + }) + .collect(); + let (server, tokens) = serve_with_documents(docs).await; + let client = connect(&server.url(), &tokens.mint("s1", "claude-code", &project)) + .await + .unwrap(); + + let (err, found) = call( + &client, + "memory_search", + json!({ "query": "grep index prefilter large repositories", "limit": 5 }), + ) + .await; + assert!(!err, "{found}"); + let size = found.to_string().len(); + assert!(size <= 4096 * 4, "{size} bytes over the default budget"); + let first = &found["documents"][0]; + let text = first["text"].as_str().unwrap(); + assert!( + text.contains("grep prefilter for large repositories"), + "{text}" + ); + assert!( + !text.contains("SESSION HEAD"), + "the excerpt is the head: {text}" + ); + assert!(text.len() < 2_000, "{} bytes in one excerpt", text.len()); + assert_eq!(first["excerpt"], true); + assert_eq!(first["length"], transcript.len()); + + // A smaller budget holds fewer documents, and says it dropped some. + let (_, tight) = call( + &client, + "memory_search", + json!({ "query": "grep prefilter", "limit": 6, "max_output_tokens": 1000 }), + ) + .await; + assert!( + tight.to_string().len() <= 4_000, + "{}", + tight.to_string().len() + ); + assert_eq!(tight["truncation"], "output_budget", "{tight}"); + assert!(tight["documents"].as_array().unwrap().len() < 6, "{tight}"); + client.cancel().await.ok(); + let _ = std::fs::remove_dir_all(&project); +} + +/// The live miss: a query no memory answers returned a branch fact and two +/// 413 notes as its top hits, each sharing one incidental word. An entry is +/// shown only when it carries the query's distinctive terms; when none does, +/// the result is empty and says so. +#[tokio::test(flavor = "multi_thread")] +async fn memory_search_does_not_present_unrelated_entries_as_matches() { + let project = temp_project("floor"); + let tokens = Arc::new(MemoryTokens::default()); + let server = serve( + ticking_memory(), + tokens.clone(), + always_on(), + Sources::default(), + ) + .await; + let client = connect(&server.url(), &tokens.mint("s1", "claude-code", &project)) + .await + .unwrap(); + for content in [ + "The repository has a local branch named '0.3.4' tracking 'origin/0.3.4'", + "prompt_too_large 413 is treated as a recoverable context-overflow signal", + "The repository uses bun for the frontend and cargo for the backend", + "Every repository change goes through a version branch", + ] { + call( + &client, + "memory_remember", + json!({ "kind": "fact", "content": content }), + ) + .await; + } + + let (err, none) = call( + &client, + "memory_search", + json!({ "query": "grep index prefilter large repositories", "limit": 5 }), + ) + .await; + assert!(!err, "{none}"); + assert_eq!(none["entries"], json!([]), "{none}"); + assert!(none["note"].is_string(), "{none}"); + + let (_, found) = call( + &client, + "memory_search", + json!({ "query": "how is a 413 prompt too large error handled" }), + ) + .await; + let contents: Vec<&str> = found["entries"] + .as_array() + .unwrap() + .iter() + .filter_map(|e| e["content"].as_str()) + .collect(); + assert_eq!( + contents, + vec!["prompt_too_large 413 is treated as a recoverable context-overflow signal"] + ); + assert!(found.get("note").is_none(), "{found}"); + client.cancel().await.ok(); + let _ = std::fs::remove_dir_all(&project); +} + +/// `max_output_tokens` is part of the schema, with the `atlas_code` tools' +/// bounds, so an agent can see it and ask for less or more. +#[test] +fn memory_search_takes_an_output_budget_like_the_code_tools() { + let wire = serde_json::to_value(tools_list()).expect("serializes"); + let search = wire["tools"] + .as_array() + .unwrap() + .iter() + .find(|t| t["name"] == "memory_search") + .expect("listed"); + let budget = &search["inputSchema"]["properties"]["max_output_tokens"]; + assert_eq!(budget["type"], "integer"); + assert_eq!(budget["minimum"], 1000); + assert_eq!(budget["maximum"], 12000); + assert!(budget["description"].as_str().unwrap().contains("4096")); + assert_eq!(search["inputSchema"]["required"], json!(["query"])); +} diff --git a/src-tauri/src/commands/memory_server/tools.rs b/src-tauri/src/commands/memory_server/tools.rs index 5630c2599..0423c1605 100644 --- a/src-tauri/src/commands/memory_server/tools.rs +++ b/src-tauri/src/commands/memory_server/tools.rs @@ -27,6 +27,7 @@ use std::sync::Arc; use atlas_memory::citation::{Citation, Validity}; use atlas_memory::record::{Entry, EntryKind, RecordStore, Source}; +use atlas_memory::retrieve::{record_query_terms, QueryTerms}; use futures::future::BoxFuture; use rmcp::handler::server::ServerHandler; use rmcp::model::{ @@ -91,6 +92,19 @@ const SEARCH_MAX_LIMIT: usize = 50; /// conversation they were meant to inform. const INDEX_DEFAULT_LIMIT: usize = 6; const INDEX_MAX_LIMIT: usize = 20; +/// The most of one document `memory_search` returns: the passage around its +/// match (about 300 tokens), never the whole text. An indexed document can be +/// a session transcript of half a megabyte. +/// The relevance floor judges a document on the same passage. +const DOC_EXCERPT_BYTES: usize = atlas_memory::retrieve::PASSAGE_BYTES; +/// The least an excerpt is cut to when a tight budget shares itself out +/// (about 200 tokens: the matching passage with a sentence either side). +/// Below it, a reply keeps fewer documents and marks the rest `truncation` +/// rather than returning fragments too short to answer anything. +const DOC_EXCERPT_MIN_BYTES: usize = 800; +/// What an empty `memory_search` says, so "nothing found" reads as an answer. +const NO_MATCH_NOTE: &str = + "nothing in shared memory or the project's documents matches this query"; /// The handoff note's JSON budget in a briefing. const HANDOFF_MAX_BYTES: usize = 4096; @@ -239,7 +253,9 @@ pub(super) fn tools() -> Vec { "kinds": { "type": "array", "items": { "type": "string", "enum": kind_names(false) }, "description": "Only these kinds (default: the durable kinds)." }, "limit": { "type": "integer", "minimum": 1, "maximum": SEARCH_MAX_LIMIT, - "description": "At most this many entries (default 10)." } + "description": "At most this many entries (default 10)." }, + "max_output_tokens": { "type": "integer", "minimum": 1000, "maximum": 12000, + "description": "Reply budget (default 4096). Documents are excerpts around the match; results past the budget are dropped, with truncation: output_budget." } }, "required": ["query"] }), @@ -387,6 +403,7 @@ struct SearchArgs { #[serde(default)] kinds: Vec, limit: Option, + max_output_tokens: Option, } #[derive(Deserialize)] @@ -626,21 +643,65 @@ pub(super) fn capped_handoff(note: &atlas_memory::handoff::HandoffNote, cap: usi } /// `result`'s JSON object with the index's `documents` added. -fn with_documents(result: CallToolResult, docs: &[IndexDoc]) -> CallToolResult { - let parsed = result - .content - .iter() - .find_map(|c| c.as_text().map(|t| t.text.clone())) - .and_then(|text| serde_json::from_str::(&text).ok()); - let Some(Value::Object(mut object)) = parsed else { - return result; +/// `max_output_tokens` as a byte budget (4 bytes a token), or the default — +/// the `atlas_code` tools' budget, read the same way. +fn output_budget(max_output_tokens: Option) -> usize { + max_output_tokens.map_or(atlas_search::DEFAULT_BUDGET_BYTES, |t| { + t.clamp(1_000, 12_000) * 4 + }) +} + +/// One indexed document as `memory_search` returns it: the passage around +/// what `terms` matched when the text is longer than `max_bytes`, marked +/// `excerpt` with the whole text's `length`. +fn document_json(d: &IndexDoc, terms: &QueryTerms, max_bytes: usize) -> Value { + let text = d.text.trim(); + match terms.excerpt(text, max_bytes) { + Some(cut) => json!({ + "title": d.title, "source": d.source, "text": cut, + "excerpt": true, "length": d.text.len(), + }), + None => json!({ "title": d.title, "source": d.source, "text": text }), + } +} + +/// `memory_search`'s reply inside `budget` bytes: entries first (the +/// record), then documents (`None`: the index was not searched), each kept +/// while it fits. A reply that dropped any says `truncation: output_budget` +/// and how many it `omitted`; an empty one carries [`NO_MATCH_NOTE`]. +fn search_reply(entries: Vec, documents: Option>, budget: usize) -> Value { + // `{"entries":[],"documents":[],"truncation":"output_budget","omitted":NN}` + // and the commas between items, reserved up front. + const FRAME: usize = 96; + let mut used = FRAME; + let mut omitted = 0usize; + let mut fit = |items: Vec| -> Vec { + let mut kept = Vec::new(); + for item in items { + let size = item.to_string().len() + 1; + if omitted == 0 && used + size <= budget { + used += size; + kept.push(item); + } else { + omitted += 1; + } + } + kept }; - let documents: Vec = docs - .iter() - .map(|d| json!({ "title": d.title, "source": d.source, "text": d.text.trim() })) - .collect(); - object.insert("documents".to_string(), Value::Array(documents)); - ok_json(Value::Object(object)) + let entries = fit(entries); + let documents = documents.map(&mut fit); + let empty = entries.is_empty() && documents.as_ref().is_none_or(Vec::is_empty); + let mut reply = json!({ "entries": entries }); + if let Some(documents) = documents { + reply["documents"] = Value::Array(documents); + } + if omitted > 0 { + reply["truncation"] = json!("output_budget"); + reply["omitted"] = json!(omitted); + } else if empty { + reply["note"] = json!(NO_MATCH_NOTE); + } + reply } /// The record entry id behind a `shared::` corpus id. @@ -787,18 +848,15 @@ impl MemoryTools { ); let handoff = run_blocking(move || { let store = shared_memory::store_for(&cwd).ok()?; - let mut note = store.last_episode(&own).ok().flatten()?; - // The previous session's facts, read now: its commits may have - // landed after it ended. - if let Some(found) = reader.stores(&cwd).find(¬e.session) { - note.apply_facts(memory_capture::session_facts(&found)); - } - Some(note) + super::briefing::read_handoff(&store, &own, &reader.stores(&cwd)) }) .await; - if let Ok(Some(note)) = handoff { - value["handoff"] = capped_handoff(¬e, HANDOFF_MAX_BYTES); - } + // Always present, as the tool description promises: `null` when no + // earlier session left a note. + value["handoff"] = match handoff { + Ok(Some(note)) => capped_handoff(¬e, HANDOFF_MAX_BYTES), + _ => Value::Null, + }; self.clocks.looked(&grant.session_id, briefing.synced_to); ok_json(value) } @@ -862,6 +920,14 @@ impl MemoryTools { /// `memory_search` over the record, plus the index when no kinds narrow /// the search to working memory. + /// + /// Two bounds the live store showed were missing. An entry is shown only + /// when it carries the query's distinctive terms ([`QueryTerms`]): RRF + /// always ranks something first, so without a floor a query nothing + /// answers got a branch fact as its top hit. (Meaning alone still finds a + /// memory — through `documents`, where the index admits a hit that stands + /// out by meaning.) And the reply has a budget, `max_output_tokens`, with + /// every document cut to the passage around its match. async fn search(&self, grant: Grant, request: CallToolRequestParams) -> CallToolResult { let args: SearchArgs = match args(&request) { Ok(a) => a, @@ -881,14 +947,21 @@ impl MemoryTools { .limit .unwrap_or(SEARCH_DEFAULT_LIMIT) .clamp(1, SEARCH_MAX_LIMIT); + let budget = output_budget(args.max_output_tokens); let (memory, cwd, query, checks) = ( self.memory.clone(), grant.cwd.clone(), args.query.clone(), self.sources.clone(), ); - let result = run_blocking(move || { + let fallback_terms = QueryTerms::new(&args.query); + let (entries, terms) = run_blocking(move || { + let terms = shared_memory::store_for(&cwd).map_or_else( + |_| QueryTerms::new(&query), + |store| record_query_terms(&store, &query), + ); let mut hits = memory.search_hits(&cwd, &query, &kinds, limit); + hits.retain(|h| terms.supports(&format!("{} {}", h.entry.key, h.entry.content))); let entries: Vec = hits.iter().map(|h| h.entry.clone()).collect(); let checked = check_entries(&checks, &cwd, &entries); // A stale memory is still found, after every one that holds. @@ -909,11 +982,11 @@ impl MemoryTools { v }) .collect(); - json!({ "entries": entries }) + (entries, terms) }) .await - .unwrap_or_else(|_| json!({ "entries": [] })); - match (&self.sources.index, args.kinds.is_empty()) { + .unwrap_or_else(|_| (Vec::new(), fallback_terms)); + let documents = match (&self.sources.index, args.kinds.is_empty()) { (Some(index), true) => { let limit = args .limit @@ -933,10 +1006,18 @@ impl MemoryTools { let docs = run_blocking(move || live_shared_docs(&memory, &cwd, docs)) .await .unwrap_or(unfiltered); - with_documents(ok_json(result), &docs) + let excerpt = DOC_EXCERPT_BYTES + .min(budget / (docs.len() + 1)) + .max(DOC_EXCERPT_MIN_BYTES); + Some( + docs.iter() + .map(|d| document_json(d, &terms, excerpt)) + .collect(), + ) } - _ => ok_json(result), - } + _ => None, + }; + ok_json(search_reply(entries, documents, budget)) } } diff --git a/src-tauri/src/commands/memory_sharing.rs b/src-tauri/src/commands/memory_sharing.rs index ee7f5643c..acc1ee696 100644 --- a/src-tauri/src/commands/memory_sharing.rs +++ b/src-tauri/src/commands/memory_sharing.rs @@ -2,9 +2,14 @@ //! //! Two per-project JSON files under `.atlas/` (atomic-written, mirroring the //! `plans.rs` / `canvas.rs` convention): -//! - `.atlas/memory-sharing.json` → `{ "enabled": bool }` (default true). -//! Gates everything: whether a session is handed the memory tool server, +//! - `.atlas/memory-sharing.json` → `{ "enabled": bool, +//! "fromExternalSessions": bool }`. `enabled` (default true) gates +//! everything: whether a session is handed the memory tool server, //! whether its deltas are captured, whether the extractor runs. +//! `fromExternalSessions` (default false) additionally lets the extractor +//! read the sessions the capture recorder imported from outside Atlas +//! (a terminal agent's own transcript), which sends their conversation +//! to the extraction model ([`MemorySharingState::extraction_reach`]). //! - `.atlas/memory-summarizer.json` → [`SummarizerPref`], the model that //! summarises the recent-session handoff `memory_briefing` serves, and //! that routes the extractor. @@ -24,6 +29,11 @@ use tauri::State; /// who never open the Memory panel still get cross-agent memory automatically. const DEFAULT_ENABLED: bool = true; +/// Default for `fromExternalSessions`: OFF. Reading a session Atlas never +/// hosted sends a conversation the user did not have in Atlas to the +/// extraction model, so it is asked for, never assumed. +const DEFAULT_FROM_EXTERNAL_SESSIONS: bool = false; + // ── Summarizer preference ──────────────────────────────────────────────────── /// Per-project handoff-summarizer preference, persisted to @@ -55,9 +65,31 @@ impl Default for SummarizerPref { // ── On-disk shape for the toggle file ──────────────────────────────────────── -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] struct SharingFile { + #[serde(default = "default_enabled")] enabled: bool, + /// Absent in files written before the setting existed: off. + #[serde(default = "default_from_external_sessions")] + from_external_sessions: bool, +} + +fn default_enabled() -> bool { + DEFAULT_ENABLED +} + +fn default_from_external_sessions() -> bool { + DEFAULT_FROM_EXTERNAL_SESSIONS +} + +impl Default for SharingFile { + fn default() -> Self { + Self { + enabled: DEFAULT_ENABLED, + from_external_sessions: DEFAULT_FROM_EXTERNAL_SESSIONS, + } + } } // ── Path helpers ───────────────────────────────────────────────────────────── @@ -86,14 +118,12 @@ fn atomic_write(path: &Path, payload: &str) -> Result<(), String> { Ok(()) } -fn read_sharing_enabled(project_path: &str) -> bool { +fn read_sharing_file(project_path: &str) -> SharingFile { let path = sharing_path(project_path); let Ok(raw) = fs::read_to_string(&path) else { - return DEFAULT_ENABLED; + return SharingFile::default(); }; - serde_json::from_str::(&raw) - .map(|f| f.enabled) - .unwrap_or(DEFAULT_ENABLED) + serde_json::from_str::(&raw).unwrap_or_default() } fn read_summarizer_pref(project_path: &str) -> SummarizerPref { @@ -106,12 +136,12 @@ fn read_summarizer_pref(project_path: &str) -> SummarizerPref { // ── Managed state ──────────────────────────────────────────────────────────── -/// The per-project toggle, cached. Registered once via `.manage()`. +/// The per-project toggles, cached. Registered once via `.manage()`. #[derive(Default)] pub struct MemorySharingState { - /// Write-through cache of the per-project enable toggle, keyed by absolute + /// Write-through cache of the per-project toggles, keyed by absolute /// project path. Avoids a file read on every send and every tool call. - toggles: Mutex>, + toggles: Mutex>, } impl MemorySharingState { @@ -119,26 +149,66 @@ impl MemorySharingState { Self::default() } - /// Whether sharing is enabled for `project_path` (cache → file → default ON). - pub fn is_enabled(&self, project_path: &str) -> bool { + /// Both toggles for `project_path` (cache → file → defaults). + fn toggles(&self, project_path: &str) -> SharingFile { if let Some(v) = self.toggles.lock().get(project_path) { return *v; } - let v = read_sharing_enabled(project_path); + let v = read_sharing_file(project_path); self.toggles.lock().insert(project_path.to_string(), v); v } + /// Whether sharing is enabled for `project_path` (default ON). + pub fn is_enabled(&self, project_path: &str) -> bool { + self.toggles(project_path).enabled + } + + /// Whether the user let extraction read `project_path`'s sessions + /// imported from outside Atlas (default OFF). The stored choice alone; + /// what extraction may read is [`Self::extraction_reach`]. + pub fn reads_external_sessions(&self, project_path: &str) -> bool { + self.toggles(project_path).from_external_sessions + } + + /// Which of the capture recorder's sessions memory extraction may read + /// in `project_path`, and so send to the extraction model. The one place + /// that decision is made: the imported (`external_jsonl`) sessions only + /// when sharing is on AND the user turned `fromExternalSessions` on; + /// otherwise only the sessions Atlas hosted. + pub fn extraction_reach(&self, project_path: &str) -> super::memory_capture::Reach { + let t = self.toggles(project_path); + if t.enabled && t.from_external_sessions { + super::memory_capture::Reach::WithImported + } else { + super::memory_capture::Reach::Hosted + } + } + /// Read the per-project summarizer preference from disk (default = raw). pub fn summarizer_pref(&self, project_path: &str) -> SummarizerPref { read_summarizer_pref(project_path) } - /// Update the toggle cache after a settings write so the next send sees it. - fn set_enabled_cache(&self, project_path: &str, enabled: bool) { - self.toggles - .lock() - .insert(project_path.to_string(), enabled); + /// Persist the `fromExternalSessions` choice for `project_path`. + pub fn set_from_external_sessions(&self, project_path: &str, on: bool) -> Result<(), String> { + self.update(project_path, |f| f.from_external_sessions = on) + } + + /// Change one toggle: read the file, apply `change`, write it back, then + /// update the cache so the next send sees it. The other toggle keeps its + /// stored value. + fn update( + &self, + project_path: &str, + change: impl FnOnce(&mut SharingFile), + ) -> Result<(), String> { + let mut file = read_sharing_file(project_path); + change(&mut file); + let payload = serde_json::to_string_pretty(&file).map_err(|e| e.to_string())?; + atomic_write(&sharing_path(project_path), &payload)?; + self.toggles.lock().insert(project_path.to_string(), file); + Ok(()) } } @@ -158,11 +228,26 @@ pub fn memory_sharing_set( enabled: bool, state: State<'_, MemorySharingState>, ) -> Result<(), String> { - let payload = - serde_json::to_string_pretty(&SharingFile { enabled }).map_err(|e| e.to_string())?; - atomic_write(&sharing_path(&project_path), &payload)?; - state.set_enabled_cache(&project_path, enabled); - Ok(()) + state.update(&project_path, |f| f.enabled = enabled) +} + +/// Whether extraction also reads the sessions imported from outside Atlas +/// (`.atlas/memory-sharing.json` → `fromExternalSessions`, default false). +#[tauri::command] +pub fn memory_from_external_sessions_get( + project_path: String, + state: State<'_, MemorySharingState>, +) -> Result { + Ok(state.reads_external_sessions(&project_path)) +} + +#[tauri::command(async)] +pub fn memory_from_external_sessions_set( + project_path: String, + enabled: bool, + state: State<'_, MemorySharingState>, +) -> Result<(), String> { + state.set_from_external_sessions(&project_path, enabled) } #[tauri::command] @@ -176,3 +261,59 @@ pub fn memory_summarizer_set(project_path: String, pref: SummarizerPref) -> Resu atomic_write(&summarizer_path(&project_path), &payload)?; Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::commands::memory_capture::Reach; + + fn project(label: &str) -> String { + let dir = + std::env::temp_dir().join(format!("atlas-sharing-{label}-{}", uuid::Uuid::new_v4())); + fs::create_dir_all(&dir).unwrap(); + dir.to_string_lossy().into_owned() + } + + /// Reading sessions imported from outside Atlas is off until asked for, + /// and only reaches extraction while sharing itself is on. + #[test] + fn imported_sessions_are_read_only_when_both_toggles_are_on() { + let p = project("reach"); + let state = MemorySharingState::new(); + assert!(state.is_enabled(&p)); + assert!(!state.reads_external_sessions(&p)); + assert_eq!(state.extraction_reach(&p), Reach::Hosted); + + state.set_from_external_sessions(&p, true).unwrap(); + assert_eq!(state.extraction_reach(&p), Reach::WithImported); + + state.update(&p, |f| f.enabled = false).unwrap(); + assert_eq!( + state.extraction_reach(&p), + Reach::Hosted, + "sharing off wins" + ); + assert!(state.reads_external_sessions(&p), "the choice is kept"); + let _ = fs::remove_dir_all(&p); + } + + /// Each toggle round-trips through the file without touching the other, + /// and a file written before the setting existed reads as off. + #[test] + fn the_toggles_round_trip_through_the_file_independently() { + let p = project("roundtrip"); + fs::create_dir_all(atlas_dir(&p)).unwrap(); + fs::write(sharing_path(&p), r#"{ "enabled": false }"#).unwrap(); + let state = MemorySharingState::new(); + assert!(!state.is_enabled(&p)); + assert!(!state.reads_external_sessions(&p)); + + state.set_from_external_sessions(&p, true).unwrap(); + let fresh = MemorySharingState::new(); + assert!(!fresh.is_enabled(&p), "enabled kept its stored value"); + assert!(fresh.reads_external_sessions(&p)); + let raw = fs::read_to_string(sharing_path(&p)).unwrap(); + assert!(raw.contains("\"fromExternalSessions\": true"), "{raw}"); + let _ = fs::remove_dir_all(&p); + } +} diff --git a/src-tauri/src/commands/org_server/tests.rs b/src-tauri/src/commands/org_server/tests.rs index 91afed774..89ef302cd 100644 --- a/src-tauri/src/commands/org_server/tests.rs +++ b/src-tauri/src/commands/org_server/tests.rs @@ -5567,6 +5567,7 @@ fn session_request(in_process: bool) -> SessionMcpRequest { http_mcp: true, ui_control: in_process, org_access: in_process, + in_process, cwd: std::path::PathBuf::from("/p"), session_id: None, } @@ -5574,23 +5575,7 @@ fn session_request(in_process: bool) -> SessionMcpRequest { /// Every entry an offer carries, as `(name, url, bearer token)`. fn entries(offer: &SessionMcpOffer) -> Vec<(String, String, String)> { - offer - .servers() - .iter() - .map(|server| { - let acp::McpServer::Http(http) = server else { - panic!("HTTP entries only") - }; - let token = http - .headers - .iter() - .find(|h| h.name == "Authorization") - .and_then(|h| h.value.strip_prefix("Bearer ")) - .expect("a bearer token") - .to_string(); - (http.name.clone(), http.url.clone(), token) - }) - .collect() + crate::commands::memory_server::offers::offered_entries(offer) } fn names(offer: &SessionMcpOffer) -> Vec { diff --git a/src-tauri/src/commands/ui_server/tests.rs b/src-tauri/src/commands/ui_server/tests.rs index 3106ed356..1d29c840d 100644 --- a/src-tauri/src/commands/ui_server/tests.rs +++ b/src-tauri/src/commands/ui_server/tests.rs @@ -412,6 +412,8 @@ fn session_request(ui_control: bool) -> SessionMcpRequest { http_mcp: true, ui_control, org_access: false, + // UI control is only ever carried by the in-process connection. + in_process: ui_control, cwd: std::path::PathBuf::from("/p"), session_id: None, } @@ -419,23 +421,7 @@ fn session_request(ui_control: bool) -> SessionMcpRequest { /// Every entry an offer carries, as `(name, url, bearer token)`. fn entries(offer: &SessionMcpOffer) -> Vec<(String, String, String)> { - offer - .servers() - .iter() - .map(|server| { - let acp::McpServer::Http(http) = server else { - panic!("HTTP entries only") - }; - let token = http - .headers - .iter() - .find(|h| h.name == "Authorization") - .and_then(|h| h.value.strip_prefix("Bearer ")) - .expect("a bearer token") - .to_string(); - (http.name.clone(), http.url.clone(), token) - }) - .collect() + crate::commands::memory_server::offers::offered_entries(offer) } fn names(offer: &SessionMcpOffer) -> Vec { diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index e966e43cd..0700536fe 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -32,20 +32,30 @@ use tauri::Manager; // that exercises the decoder rather than by a const that only proves a patch // still applies. -/// `atlas mcp-bridge `: run the stdio bridge (ADR-0019) instead of the -/// app, the session token from `ATLAS_MCP_TOKEN`. Returns whether it ran. +/// `atlas mcp-bridge `: run the stdio bridge (ADR-0019) +/// instead of the app, the session token read from the file (ADR-0020: never +/// on the command line, never in the environment an agent may print there). +/// Returns whether it ran. pub fn run_bridge_if_asked() -> bool { - let mut args = std::env::args().skip(1); - if args.next().as_deref() != Some(commands::memory_server::BRIDGE_ARG) { + let mut args = std::env::args_os().skip(1); + if args.next().and_then(|a| a.into_string().ok()).as_deref() + != Some(commands::memory_server::BRIDGE_ARG) + { return false; } - let (Some(url), Ok(token)) = ( - args.next(), - std::env::var(commands::memory_server::BRIDGE_TOKEN_ENV), - ) else { - eprintln!("usage: ATLAS_MCP_TOKEN=... atlas mcp-bridge "); + let (Some(url), Some(token_file)) = + (args.next().and_then(|a| a.into_string().ok()), args.next()) + else { + eprintln!("usage: atlas mcp-bridge "); std::process::exit(2); }; + let token = match commands::memory_bridge::read_token_file(std::path::Path::new(&token_file)) { + Ok(token) => token, + Err(e) => { + eprintln!("mcp-bridge: {e:#}"); + std::process::exit(1); + } + }; let runtime = match tokio::runtime::Builder::new_current_thread() .enable_all() .build() @@ -922,6 +932,8 @@ pub fn run() { commands::memory_policy::memory_policy_update, commands::memory_sharing::memory_sharing_get, commands::memory_sharing::memory_sharing_set, + commands::memory_sharing::memory_from_external_sessions_get, + commands::memory_sharing::memory_from_external_sessions_set, commands::memory_sharing::memory_summarizer_get, commands::memory_sharing::memory_summarizer_set, commands::shared_memory::memory_get_state, diff --git a/src-tauri/src/main.rs b/src-tauri/src/main.rs index d9c01a1a0..f4e0f1699 100644 --- a/src-tauri/src/main.rs +++ b/src-tauri/src/main.rs @@ -8,8 +8,9 @@ static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; fn main() { - // `atlas mcp-bridge ` is the stdio bridge for agents without HTTP - // MCP (ADR-0019), not the app. + // `atlas mcp-bridge ` is the stdio bridge every ACP + // agent reaches Atlas's tool servers through (ADR-0019, ADR-0020), not + // the app. if atlas_lib::run_bridge_if_asked() { return; } diff --git a/src/dev/mock-backend/fixtures/memory.ts b/src/dev/mock-backend/fixtures/memory.ts index 59a76fa96..f56e576bc 100644 --- a/src/dev/mock-backend/fixtures/memory.ts +++ b/src/dev/mock-backend/fixtures/memory.ts @@ -769,6 +769,9 @@ const sharingEnabled = new Map([ [ALL_PROJECTS[2].path, false], ]); +/** Reading sessions imported from outside Atlas: off by default, as in Rust. */ +const fromExternalSessions = new Map(); + /** The BYOK summariser, so the provider/model picker opens on a selection * rather than on the default `raw` (which hides the picker entirely). */ const summarizers = new Map([ @@ -1284,6 +1287,8 @@ export interface MemoryResponses { memory_policy_update: Unit; memory_sharing_get: boolean; memory_sharing_set: Unit; + memory_from_external_sessions_get: boolean; + memory_from_external_sessions_set: Unit; memory_summarizer_get: SummarizerPref; memory_summarizer_set: Unit; memory_get_state: SharedState; @@ -1375,6 +1380,12 @@ export const memoryHandlers: TypedHandlers = { sharingEnabled.set(String(projectPath), Boolean(enabled)); return null; }, + memory_from_external_sessions_get: ({ projectPath }): boolean => + fromExternalSessions.get(String(projectPath)) ?? false, + memory_from_external_sessions_set: ({ projectPath, enabled }): null => { + fromExternalSessions.set(String(projectPath), Boolean(enabled)); + return null; + }, memory_summarizer_get: ({ projectPath }): SummarizerPref => summarizers.get(String(projectPath)) ?? DEFAULT_SUMMARIZER, memory_summarizer_set: ({ projectPath, pref }): null => { diff --git a/src/dev/mock-backend/scenarios/northwind-misc.ts b/src/dev/mock-backend/scenarios/northwind-misc.ts index 7838240ab..4123b2874 100644 --- a/src/dev/mock-backend/scenarios/northwind-misc.ts +++ b/src/dev/mock-backend/scenarios/northwind-misc.ts @@ -777,6 +777,7 @@ const EVENT_KIND: Record = { const entryEdits = new Map(); const forgotten = new Set(); let sharing = true; +let fromExternalSessions = false; function entries(): MemoryEntry[] { return memory() @@ -1449,6 +1450,11 @@ export const northwindMiscCommands: Partial> = { sharing = Boolean(enabled); return null; }, + memory_from_external_sessions_get: (): boolean => fromExternalSessions, + memory_from_external_sessions_set: ({ enabled }): null => { + fromExternalSessions = Boolean(enabled); + return null; + }, memory_summarizer_get: (): SummarizerPref => ({ mode: "raw", provider: "", model: "" }), memory_get_state: ({ projectPath }): SharedState => isProject(projectPath) diff --git a/src/features/memory/components/memory-sharing-controls.test.tsx b/src/features/memory/components/memory-sharing-controls.test.tsx new file mode 100644 index 000000000..1082796c0 --- /dev/null +++ b/src/features/memory/components/memory-sharing-controls.test.tsx @@ -0,0 +1,82 @@ +// @vitest-environment happy-dom +// The Memory panel's sharing popover: the "sessions outside Atlas" switch is +// off by default, persists through its own command, and is inert while Shared +// is off (nothing is extracted then, so it would promise nothing). +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +const invoke = vi.fn(); +vi.mock("@tauri-apps/api/core", () => ({ invoke: (...args: unknown[]) => invoke(...args) })); +vi.mock("@tauri-apps/api/event", () => ({ listen: vi.fn(async () => () => {}) })); + +import { cleanup, render, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { MemorySharingControls } from "./memory-sharing-controls"; +import { useMemorySharingStore } from "../stores/memory-sharing-store"; +import { + EXTERNAL_SESSIONS_LABEL, + EXTERNAL_SESSIONS_NEEDS_SHARING, +} from "../lib/memory-sharing-copy"; + +const PROJECT = "/work/acme"; + +function backend({ sharing, external }: { sharing: boolean; external: boolean }) { + invoke.mockImplementation(async (cmd: string) => { + switch (cmd) { + case "memory_sharing_get": + return sharing; + case "memory_from_external_sessions_get": + return external; + case "memory_summarizer_get": + return { mode: "raw", provider: "", model: "" }; + case "byok_list": + return {}; + default: + return null; + } + }); +} + +async function openPopover() { + render(); + await waitFor(() => expect(useMemorySharingStore.getState().loaded).toBe(true)); + await userEvent.click(screen.getByLabelText("Memory sharing settings")); + return screen.findByRole("switch", { name: EXTERNAL_SESSIONS_LABEL }); +} + +beforeEach(() => { + invoke.mockReset(); + useMemorySharingStore.setState({ + projectPath: null, + enabled: true, + fromExternalSessions: false, + loaded: false, + }); +}); +afterEach(cleanup); + +describe("the sessions-outside-Atlas switch", () => { + it("is off by default and turning it on persists for the project", async () => { + backend({ sharing: true, external: false }); + const toggle = await openPopover(); + expect(toggle.getAttribute("aria-checked")).toBe("false"); + expect(toggle.hasAttribute("disabled")).toBe(false); + + await userEvent.click(toggle); + expect(invoke).toHaveBeenCalledWith("memory_from_external_sessions_set", { + projectPath: PROJECT, + enabled: true, + }); + expect(toggle.getAttribute("aria-checked")).toBe("true"); + }); + + it("shows the stored choice but is disabled while Shared is off", async () => { + backend({ sharing: false, external: true }); + const toggle = await openPopover(); + expect(toggle.getAttribute("aria-checked")).toBe("true"); + expect(toggle.hasAttribute("disabled")).toBe(true); + expect(screen.getByText(EXTERNAL_SESSIONS_NEEDS_SHARING)).toBeTruthy(); + + await userEvent.click(toggle); + expect(invoke).not.toHaveBeenCalledWith("memory_from_external_sessions_set", expect.anything()); + }); +}); diff --git a/src/features/memory/components/memory-sharing-controls.tsx b/src/features/memory/components/memory-sharing-controls.tsx index 79f550d4e..6a5015e3d 100644 --- a/src/features/memory/components/memory-sharing-controls.tsx +++ b/src/features/memory/components/memory-sharing-controls.tsx @@ -3,7 +3,9 @@ // Two affordances, kept to the monochrome/hairline house style (see Atlas // Design Principles): a Shared toggle pill (white when on) and a settings // popover holding the handoff-summarizer mode selector (Raw / Provider / -// Local-disabled) plus the reused ProviderModelSelector when mode === provider. +// Local-disabled), the reused ProviderModelSelector when mode === provider, +// and the "sessions outside Atlas" switch (off by default; inert while Shared +// is off, since nothing is extracted then). import { useEffect, useMemo } from "react"; import { Popover } from "@base-ui/react/popover"; @@ -11,16 +13,26 @@ import { Share2, SlidersHorizontal, FileText, Server, Cpu, Check } from "lucide- import { cn } from "@/lib/utils"; import { Hint } from "@/ui/tooltip"; import { ProviderModelSelector } from "./provider-pickers"; +import { Toggle } from "@/features/settings/components/settings-controls"; import { useByokStore } from "@/features/settings/stores/byok-store"; import { CHAT_PROVIDERS } from "@/features/settings/lib/providers"; import { useMemorySharingStore } from "../stores/memory-sharing-store"; import type { SummarizerMode } from "../lib/memory-sharing-api"; -import { EXTRACTION_NOTE, HANDOFF_HINT, RAW_HINT } from "../lib/memory-sharing-copy"; +import { + EXTERNAL_SESSIONS_HINT, + EXTERNAL_SESSIONS_LABEL, + EXTERNAL_SESSIONS_NEEDS_SHARING, + EXTRACTION_NOTE, + HANDOFF_HINT, + RAW_HINT, +} from "../lib/memory-sharing-copy"; export function MemorySharingControls({ projectPath }: { projectPath: string | null }) { const enabled = useMemorySharingStore.use.enabled(); const pref = useMemorySharingStore.use.pref(); - const { load, setEnabled, setPref } = useMemorySharingStore.use.actions(); + const fromExternalSessions = useMemorySharingStore.use.fromExternalSessions(); + const { load, setEnabled, setFromExternalSessions, setPref } = + useMemorySharingStore.use.actions(); const byokKeys = useByokStore.use.keys(); const byokLoaded = useByokStore.use.loaded(); @@ -70,7 +82,7 @@ export function MemorySharingControls({ projectPath }: { projectPath: string | n {/* Summarizer settings popover */} - + {EXTRACTION_NOTE}

+ +
+
+
+

{EXTERNAL_SESSIONS_LABEL}

+

+ {EXTERNAL_SESSIONS_HINT} +

+ {!enabled && ( +

+ {EXTERNAL_SESSIONS_NEEDS_SHARING} +

+ )} +
+ void setFromExternalSessions(on)} + /> +
+
diff --git a/src/features/memory/lib/memory-sharing-api.ts b/src/features/memory/lib/memory-sharing-api.ts index f684bab83..17da30c2c 100644 --- a/src/features/memory/lib/memory-sharing-api.ts +++ b/src/features/memory/lib/memory-sharing-api.ts @@ -1,4 +1,5 @@ -// Shared Cross-Agent Memory — TS bindings for the per-project toggle + handoff +// Shared Cross-Agent Memory — TS bindings for the per-project toggles (sharing, +// and whether extraction reads sessions imported from outside Atlas) + handoff // summarizer preference. Agents pull memory through the `atlas_memory` tools // (ADR-0010); nothing is injected. These commands only read/write the two // per-project settings files. Mirrors the plain-invoke pattern in @@ -21,6 +22,11 @@ export const memorySharing = { getEnabled: (projectPath: string) => invoke("memory_sharing_get", { projectPath }), setEnabled: (projectPath: string, enabled: boolean) => invoke("memory_sharing_set", { projectPath, enabled }), + /** Whether extraction also reads sessions imported from outside Atlas (default false). */ + getFromExternalSessions: (projectPath: string) => + invoke("memory_from_external_sessions_get", { projectPath }), + setFromExternalSessions: (projectPath: string, enabled: boolean) => + invoke("memory_from_external_sessions_set", { projectPath, enabled }), getSummarizer: (projectPath: string) => invoke("memory_summarizer_get", { projectPath }), setSummarizer: (projectPath: string, pref: SummarizerPref) => diff --git a/src/features/memory/lib/memory-sharing-copy.test.ts b/src/features/memory/lib/memory-sharing-copy.test.ts index 3cfe2ee5b..6d8a6dddd 100644 --- a/src/features/memory/lib/memory-sharing-copy.test.ts +++ b/src/features/memory/lib/memory-sharing-copy.test.ts @@ -1,5 +1,10 @@ import { describe, expect, it } from "vitest"; -import { EXTRACTION_NOTE, HANDOFF_HINT, RAW_HINT } from "./memory-sharing-copy"; +import { + EXTERNAL_SESSIONS_HINT, + EXTRACTION_NOTE, + HANDOFF_HINT, + RAW_HINT, +} from "./memory-sharing-copy"; describe("the summarizer popover copy", () => { it("never claims context is injected", () => { @@ -13,3 +18,11 @@ describe("the summarizer popover copy", () => { expect(RAW_HINT).toMatch(/No model call for the handoff/); }); }); + +describe("the sessions-outside-Atlas copy", () => { + it("says what is read, that it is sent to the model, and that it starts off", () => { + expect(EXTERNAL_SESSIONS_HINT).toMatch(/imported from your terminal/); + expect(EXTERNAL_SESSIONS_HINT).toMatch(/sent to the extraction model/); + expect(EXTERNAL_SESSIONS_HINT).toMatch(/Off by default/); + }); +}); diff --git a/src/features/memory/lib/memory-sharing-copy.ts b/src/features/memory/lib/memory-sharing-copy.ts index 58a333e63..2e5838221 100644 --- a/src/features/memory/lib/memory-sharing-copy.ts +++ b/src/features/memory/lib/memory-sharing-copy.ts @@ -13,3 +13,13 @@ export const RAW_HINT = /** Always shown: what extraction does with the session, whatever the mode. */ export const EXTRACTION_NOTE = "While Shared is on and you're signed in, Atlas also sends a redacted excerpt of long sessions to the Atlas model to extract decisions and facts. Choose Provider to use your own key instead."; + +/** The switch that lets extraction read sessions Atlas did not host. */ +export const EXTERNAL_SESSIONS_LABEL = "Learn from sessions outside Atlas"; + +/** What that switch does, said plainly: what is read and where it goes. */ +export const EXTERNAL_SESSIONS_HINT = + "Also extracts memories from coding-agent sessions Atlas imported from your terminal. Their transcripts are sent to the extraction model, the same as in-app sessions. Off by default."; + +/** Shown instead while Shared is off, when the switch can't do anything. */ +export const EXTERNAL_SESSIONS_NEEDS_SHARING = "Turn on Shared to use this."; diff --git a/src/features/memory/stores/memory-sharing-store.ts b/src/features/memory/stores/memory-sharing-store.ts index 4e58c14b2..bc7264f36 100644 --- a/src/features/memory/stores/memory-sharing-store.ts +++ b/src/features/memory/stores/memory-sharing-store.ts @@ -12,11 +12,14 @@ const DEFAULT_PREF: SummarizerPref = { mode: "raw", provider: "", model: "" }; interface MemorySharingState { projectPath: string | null; enabled: boolean; + /** Extraction also reads sessions imported from outside Atlas. */ + fromExternalSessions: boolean; pref: SummarizerPref; loaded: boolean; actions: { load: (projectPath: string) => Promise; setEnabled: (enabled: boolean) => Promise; + setFromExternalSessions: (on: boolean) => Promise; setPref: (pref: SummarizerPref) => Promise; }; } @@ -25,19 +28,21 @@ export const useMemorySharingStore = createSelectors( create((set, get) => ({ projectPath: null, enabled: true, // optimistic default (matches Rust DEFAULT_ENABLED) + fromExternalSessions: false, // matches Rust DEFAULT_FROM_EXTERNAL_SESSIONS pref: DEFAULT_PREF, loaded: false, actions: { load: async (projectPath) => { set({ projectPath, loaded: false }); try { - const [enabled, pref] = await Promise.all([ + const [enabled, fromExternalSessions, pref] = await Promise.all([ memorySharing.getEnabled(projectPath), + memorySharing.getFromExternalSessions(projectPath), memorySharing.getSummarizer(projectPath), ]); // Ignore a stale response if the project changed mid-flight. if (get().projectPath !== projectPath) return; - set({ enabled, pref, loaded: true }); + set({ enabled, fromExternalSessions, pref, loaded: true }); } catch (err) { console.error("memorySharing.load failed", err); set({ loaded: true }); @@ -57,6 +62,19 @@ export const useMemorySharingStore = createSelectors( } }, + setFromExternalSessions: async (on) => { + const projectPath = get().projectPath; + if (!projectPath) return; + const prev = get().fromExternalSessions; + set({ fromExternalSessions: on }); // optimistic + try { + await memorySharing.setFromExternalSessions(projectPath, on); + } catch (err) { + console.error("memorySharing.setFromExternalSessions failed", err); + set({ fromExternalSessions: prev }); // revert + } + }, + setPref: async (pref) => { const projectPath = get().projectPath; if (!projectPath) return; diff --git a/src/features/settings/components/settings-controls.tsx b/src/features/settings/components/settings-controls.tsx index d70a38a88..ac4cc5561 100644 --- a/src/features/settings/components/settings-controls.tsx +++ b/src/features/settings/components/settings-controls.tsx @@ -40,6 +40,8 @@ export function Toggle({ checked, onChange, disabled = false, + label, + className, }: { defaultChecked?: boolean; checked?: boolean; @@ -47,6 +49,11 @@ export function Toggle({ /** For a sub-setting whose parent is off — dimmed and inert, but still * showing its own stored value rather than lying about it. */ disabled?: boolean; + /** Accessible name, for a switch whose visible label is not a