bingshuo/grok-build-upstream-mirror
Watch
1
0
Fork
You've already forked grok-build-upstream-mirror
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Publish harness and TUI open-source

Browse source 
initial sync from the monorepo
This commit is contained in:
grokkybara[bot] 2026-07-16 06:46:02 +01:00
commit c68e39f604
2734 changed files with 1437016 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

50
crates/codegen/xai-grok-memory/Cargo.toml Normal file
View file

@ -0,0 +1,50 @@
[package]
license = "Apache-2.0"
name = "xai-grok-memory"
version = "0.1.0"
edition.workspace = true
[features]
[dependencies]
anyhow = { workspace = true }
arc-swap = { workspace = true }
async-trait = { workspace = true }
blake3 = { workspace = true }
chrono = { workspace = true }
dunce = { workspace = true }
flate2 = { workspace = true }
git2 = { version = "0.20", default-features = false, features = [
"vendored-libgit2",
] }
notify = { workspace = true }
reqwest = { workspace = true }
reqwest-middleware = { workspace = true }
# rusqlite 0.37 + bundled = self-contained SQLite (>= 3.50.2) with FTS5.
rusqlite = { version = "0.37", features = ["bundled"] }
serde_json = { workspace = true }
# sqlite-vec: vec0 virtual table for KNN vector search.
sqlite-vec = "=0.1.7-alpha.2"
tar = { workspace = true }
tokio = { workspace = true, features = ["macros", "rt", "rt-multi-thread", "time"] }
tracing = { workspace = true }
xai-grok-auth = { workspace = true }
xai-grok-config-types = { workspace = true }
xai-grok-http = { workspace = true }
xai-grok-telemetry = { workspace = true }
xai-grok-tools = { workspace = true }
xai-grok-version = { workspace = true }
xai-sqlite-journal = { workspace = true }
[target.'cfg(unix)'.dependencies]
nix = { workspace = true }
[target.'cfg(windows)'.dependencies]
windows = { workspace = true }
[dev-dependencies]
filetime = { workspace = true }
tempfile = { workspace = true }
[lints]
workspace = true

116
crates/codegen/xai-grok-memory/src/archive.rs Normal file
View file

@ -0,0 +1,116 @@
//! Build a `memory.tar.gz` archive containing session logs and MEMORY.md files.
//!
//! The archive is uploaded to GCS at session finalize time. The reconstruct
//! pipeline injects these into the Docker image for full replay fidelity.
use anyhow::{Context, Result};
use super::MemoryStorage;
/// Build a `memory.tar.gz` archive with session logs and MEMORY.md files.
pub fn build_memory_archive(storage: &MemoryStorage) -> Result<Vec<u8>> {
use flate2::Compression;
use flate2::write::GzEncoder;
let buf = Vec::new();
let enc = GzEncoder::new(buf, Compression::default());
let mut ar = tar::Builder::new(enc);
// Session logs
let sessions_dir = storage.workspace_dir().join("sessions");
if sessions_dir.is_dir() {
for entry in std::fs::read_dir(&sessions_dir)
.context("read sessions dir")?
.flatten()
{
let path = entry.path();
if path.extension().and_then(|e| e.to_str()) == Some("md") {
let name = format!("workspace/sessions/{}", entry.file_name().to_string_lossy());
ar.append_path_with_name(&path, &name)
.with_context(|| format!("archive {name}"))?;
}
}
}
// MEMORY.md files
let global_mem = storage.global_memory_file();
if global_mem.is_file() {
ar.append_path_with_name(&global_mem, "global/MEMORY.md")
.context("archive global MEMORY.md")?;
}
let workspace_mem = storage.workspace_memory_file();
if workspace_mem.is_file() {
let ws_dir_name = storage
.workspace_dir()
.file_name()
.and_then(|n| n.to_str())
.unwrap_or("workspace");
let archive_path = format!("{ws_dir_name}/MEMORY.md");
ar.append_path_with_name(&workspace_mem, &archive_path)
.context("archive workspace MEMORY.md")?;
}
let enc = ar.into_inner().context("finalize tar")?;
enc.finish().context("compress tar.gz")
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
fn test_storage(tmp: &TempDir) -> MemoryStorage {
let global = tmp.path().join("memory");
let workspace = global.join("test_ws");
MemoryStorage::with_paths(global, workspace)
}
#[test]
fn test_build_empty_archive() {
let tmp = TempDir::new().unwrap();
let storage = test_storage(&tmp);
let archive = build_memory_archive(&storage).unwrap();
assert!(!archive.is_empty());
}
#[test]
fn test_build_archive_with_files() {
let tmp = TempDir::new().unwrap();
let storage = test_storage(&tmp);
storage.ensure_initialized().unwrap();
storage
.write_daily_log("2026-03-09", "test", "sess12345678", "# Test", false)
.unwrap();
let archive = build_memory_archive(&storage).unwrap();
assert!(archive.len() > 100);
}
#[test]
fn test_build_archive_includes_memory_md() {
let tmp = TempDir::new().unwrap();
let storage = test_storage(&tmp);
storage.ensure_initialized().unwrap();
std::fs::write(storage.global_memory_file(), "# Global Memory").unwrap();
std::fs::write(storage.workspace_memory_file(), "# Workspace Memory").unwrap();
let archive = build_memory_archive(&storage).unwrap();
let entries = tar_entry_names(&archive);
assert!(entries.contains(&"global/MEMORY.md".to_string()));
assert!(entries.contains(&"test_ws/MEMORY.md".to_string()));
}
fn tar_entry_names(gz_bytes: &[u8]) -> Vec<String> {
use flate2::read::GzDecoder;
let decoder = GzDecoder::new(gz_bytes);
let mut archive = tar::Archive::new(decoder);
archive
.entries()
.unwrap()
.filter_map(|e| e.ok())
.map(|e| e.path().unwrap().display().to_string())
.collect()
}
}

1415
crates/codegen/xai-grok-memory/src/backend.rs Normal file
View file

File diff suppressed because it is too large Load diff

367
crates/codegen/xai-grok-memory/src/chunker.rs Normal file
View file

@ -0,0 +1,367 @@
//! Markdown-aware semantic chunking.
//!
//! Splits markdown content into chunks suitable for embedding and search.
//! Chunks respect markdown structure (headers, paragraphs, code blocks)
//! and include ancestor headers for self-containment.
//!
//! Character counts are used as a proxy for token counts (chars / 4 ≈ tokens).
use xai_grok_config_types::MemoryIndexConfig;
/// A chunk of text extracted from a memory file.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Chunk {
/// The chunk text, including ancestor header context.
pub text: String,
/// 0-based start line in the source file.
pub start_line: usize,
/// 0-based end line (exclusive) in the source file.
pub end_line: usize,
}
/// Compute a blake3 hash of the chunk text, returned as a hex string.
pub fn chunk_hash(text: &str) -> String {
blake3::hash(text.as_bytes()).to_hex().to_string()
}
/// Split markdown content into chunks, respecting structure.
///
/// Strategy:
/// 1. Split on `##` headers — each section is a candidate chunk
/// 2. If a section exceeds `max_chunk_chars`, split on paragraph boundaries (`\n\n`)
/// 3. If a paragraph still exceeds `max_chunk_chars`, split on line boundaries
/// 4. Continuation chunks are prefixed with ancestor header context
///
/// When a section is split into multiple sub-chunks, each continuation chunk
/// is prefixed with the last `chunk_overlap_chars` of the previous chunk for
/// embedding continuity, plus ancestor header context.
pub fn chunk_markdown(content: &str, config: &MemoryIndexConfig) -> Vec<Chunk> {
if content.is_empty() {
return vec![];
}
let max_chars = config.max_chunk_chars;
let lines: Vec<&str> = content.lines().collect();
if lines.is_empty() {
return vec![];
}
// If the entire content fits in one chunk, return it directly.
if content.len() <= max_chars {
return vec![Chunk {
text: content.to_string(),
start_line: 0,
end_line: lines.len(),
}];
}
// Split into sections by ## headers
let sections = split_by_headers(&lines);
let mut chunks = Vec::new();
for section in &sections {
let section_text = section.lines.join("\n");
if section_text.len() <= max_chars {
chunks.push(Chunk {
text: add_header_context(&section.header_context, &section_text),
start_line: section.start_line,
end_line: section.start_line + section.lines.len(),
});
} else {
// Section too large — split on paragraph boundaries
let sub_chunks =
split_section_by_paragraphs(section, max_chars, config.chunk_overlap_chars);
chunks.extend(sub_chunks);
}
}
chunks
}
/// A section of the document delimited by headers.
struct Section<'a> {
/// The lines in this section (including the header line itself).
lines: Vec<&'a str>,
/// 0-based start line index in the original document.
start_line: usize,
/// Ancestor header context (e.g., `"## Architecture > ### Design"`).
header_context: String,
}
/// Split lines into sections by `##` (or deeper) headers.
fn split_by_headers<'a>(lines: &[&'a str]) -> Vec<Section<'a>> {
let mut sections: Vec<Section<'a>> = Vec::new();
let mut current_lines: Vec<&'a str> = Vec::new();
let mut current_start = 0;
let mut header_stack: Vec<(usize, String)> = Vec::new(); // (level, text)
for (i, &line) in lines.iter().enumerate() {
if let Some(level) = header_level(line) {
// Flush previous section
if !current_lines.is_empty() {
sections.push(Section {
lines: std::mem::take(&mut current_lines),
start_line: current_start,
header_context: format_header_context(&header_stack),
});
}
current_start = i;
// Update header stack: pop headers at same or deeper level
while header_stack.last().is_some_and(|(l, _)| *l >= level) {
header_stack.pop();
}
header_stack.push((level, line.to_string()));
}
current_lines.push(line);
}
// Flush final section
if !current_lines.is_empty() {
sections.push(Section {
lines: current_lines,
start_line: current_start,
header_context: format_header_context(&header_stack),
});
}
sections
}
/// Split a large section into sub-chunks by paragraph boundaries (`\n\n`).
/// Continuation chunks are prefixed with the last `overlap_chars` of the
/// previous chunk for embedding continuity.
fn split_section_by_paragraphs(
section: &Section<'_>,
max_chars: usize,
overlap_chars: usize,
) -> Vec<Chunk> {
let mut chunks = Vec::new();
let mut current_text = String::new();
let mut current_start = section.start_line;
let mut line_offset = 0;
for (i, &line) in section.lines.iter().enumerate() {
let is_blank = line.trim().is_empty();
// Paragraph boundary: blank line AND accumulated text is non-empty
if is_blank && !current_text.is_empty() && current_text.len() + line.len() > max_chars {
// Flush current chunk
let flushed = current_text.trim().to_string();
chunks.push(Chunk {
text: add_header_context(&section.header_context, &flushed),
start_line: current_start,
end_line: section.start_line + i,
});
// Apply overlap: start next chunk with tail of previous
current_text = if overlap_chars > 0 {
let tail: String = flushed
.chars()
.rev()
.take(overlap_chars)
.collect::<Vec<_>>()
.into_iter()
.rev()
.collect();
tail
} else {
String::new()
};
current_start = section.start_line + i + 1;
line_offset = i + 1;
continue;
}
if !current_text.is_empty() {
current_text.push('\n');
}
current_text.push_str(line);
// If single line pushes us over max, flush what we have
if current_text.len() > max_chars && i > line_offset {
// Split at the previous line
let split_at = current_text.rfind('\n').unwrap_or(current_text.len());
let (keep, remainder) = current_text.split_at(split_at);
chunks.push(Chunk {
text: add_header_context(&section.header_context, keep.trim()),
start_line: current_start,
end_line: section.start_line + i,
});
current_text = remainder.trim_start_matches('\n').to_string();
current_start = section.start_line + i;
line_offset = i;
}
}
// Flush remaining
if !current_text.trim().is_empty() {
chunks.push(Chunk {
text: add_header_context(&section.header_context, current_text.trim()),
start_line: current_start,
end_line: section.start_line + section.lines.len(),
});
}
chunks
}
/// Detect markdown header level (1 for `#`, 2 for `##`, etc.). Returns `None` if not a header.
pub(crate) fn header_level(line: &str) -> Option<usize> {
let trimmed = line.trim_start();
if !trimmed.starts_with('#') {
return None;
}
let level = trimmed.chars().take_while(|&c| c == '#').count();
// Must be followed by a space or end of line to be a valid header
let rest = &trimmed[level..];
if rest.is_empty() || rest.starts_with(' ') {
Some(level)
} else {
None
}
}
/// Format header stack into a context string like `"## Section > ### Subsection"`.
fn format_header_context(stack: &[(usize, String)]) -> String {
if stack.len() <= 1 {
return String::new();
}
// Skip the last entry (it's the current section's own header)
stack[..stack.len() - 1]
.iter()
.map(|(_, text)| text.trim().to_string())
.collect::<Vec<_>>()
.join(" > ")
}
/// Prepend ancestor header context to chunk text (if non-empty).
fn add_header_context(context: &str, text: &str) -> String {
if context.is_empty() {
text.to_string()
} else {
format!("[Context: {context}]\n\n{text}")
}
}
#[cfg(test)]
mod tests {
use super::*;
fn default_config() -> MemoryIndexConfig {
MemoryIndexConfig::default()
}
#[test]
fn test_chunk_hash_deterministic() {
let h1 = chunk_hash("hello world");
let h2 = chunk_hash("hello world");
assert_eq!(h1, h2);
assert_eq!(h1.len(), 64); // blake3 hex = 64 chars
}
#[test]
fn test_chunk_hash_different_inputs() {
assert_ne!(chunk_hash("hello"), chunk_hash("world"));
}
#[test]
fn test_chunk_empty_content() {
let chunks = chunk_markdown("", &default_config());
assert!(chunks.is_empty());
}
#[test]
fn test_chunk_small_content_single_chunk() {
let content = "# Title\n\nSome text here.";
let chunks = chunk_markdown(content, &default_config());
assert_eq!(chunks.len(), 1);
assert_eq!(chunks[0].text, content);
assert_eq!(chunks[0].start_line, 0);
assert_eq!(chunks[0].end_line, 3);
}
#[test]
fn test_chunk_splits_on_headers() {
let content = "## Section 1\n\nContent for section 1 goes here with enough text to matter.\n\n\
## Section 2\n\nContent for section 2 is also significant enough to be a chunk.";
let config = MemoryIndexConfig {
max_chunk_chars: 80,
chunk_overlap_chars: 0,
};
let chunks = chunk_markdown(content, &config);
assert!(
chunks.len() >= 2,
"should split into at least 2 chunks, got {}",
chunks.len()
);
assert!(chunks[0].text.contains("Section 1"));
assert!(chunks.last().unwrap().text.contains("Section 2"));
}
#[test]
fn test_chunk_header_context_for_subsections() {
let content = "## Parent\n\nIntro.\n\n### Child\n\nChild content that is long enough to be its own chunk definitely.";
let config = MemoryIndexConfig {
max_chunk_chars: 60,
chunk_overlap_chars: 0,
};
let chunks = chunk_markdown(content, &config);
// The child section chunk should have parent context
let child_chunk = chunks.iter().find(|c| c.text.contains("Child content"));
assert!(child_chunk.is_some(), "should have a child chunk");
assert!(
child_chunk.unwrap().text.contains("[Context: ## Parent]"),
"child chunk should have parent header context, got: {}",
child_chunk.unwrap().text
);
}
#[test]
fn test_chunk_large_section_splits_on_paragraphs() {
let para1 = "A".repeat(100);
let para2 = "B".repeat(100);
let content = format!("## Big Section\n\n{para1}\n\n{para2}");
let config = MemoryIndexConfig {
max_chunk_chars: 150,
chunk_overlap_chars: 0,
};
let chunks = chunk_markdown(&content, &config);
assert!(
chunks.len() >= 2,
"should split large section, got {} chunks",
chunks.len()
);
}
#[test]
fn test_header_level_detection() {
assert_eq!(header_level("# Title"), Some(1));
assert_eq!(header_level("## Section"), Some(2));
assert_eq!(header_level("### Subsection"), Some(3));
assert_eq!(header_level("#hashtag"), None); // no space after #
assert_eq!(header_level("not a header"), None);
assert_eq!(header_level(""), None);
assert_eq!(header_level("##"), Some(2)); // header with no text
}
#[test]
fn test_chunk_line_numbers() {
let content = "line 0\nline 1\nline 2\nline 3\nline 4";
let chunks = chunk_markdown(content, &default_config());
assert_eq!(chunks.len(), 1);
assert_eq!(chunks[0].start_line, 0);
assert_eq!(chunks[0].end_line, 5);
}
#[test]
fn test_chunk_preserves_code_blocks() {
let content =
"## Code\n\n```rust\nfn main() {\n println!(\"hello\");\n}\n```\n\nSome text.";
let chunks = chunk_markdown(content, &default_config());
assert_eq!(chunks.len(), 1);
assert!(chunks[0].text.contains("```rust"));
assert!(chunks[0].text.contains("fn main()"));
}
}

1471
crates/codegen/xai-grok-memory/src/dream.rs Normal file
View file

File diff suppressed because it is too large Load diff

488
crates/codegen/xai-grok-memory/src/dream_lock.rs Normal file
View file

@ -0,0 +1,488 @@
//! Dream lock file and session counting infrastructure.
//!
//! Coordination primitives for background memory consolidation ("dream"):
//! - [`DreamLock`]: PID-based lock file with mtime tracking
//! - [`sessions_since`]: counts session files modified after a given timestamp
use std::fs;
use std::io;
use std::path::{Path, PathBuf};
use std::time::SystemTime;
const LOCK_FILE_NAME: &str = ".dream-lock";
/// Whether a process with the given PID is alive.
///
/// Local copy kept dependency-free of `crate::util` so the memory subsystem
/// can be extracted into its own crate. Mirrors `crate::util::is_process_alive`.
#[cfg(unix)]
fn is_process_alive(pid: u32) -> bool {
use nix::errno::Errno;
use nix::sys::signal::kill;
use nix::unistd::Pid;
// Signal 0 probes existence; EPERM means alive under a different UID.
match kill(Pid::from_raw(pid as i32), None) {
Ok(()) => true,
Err(Errno::ESRCH) => false,
Err(_) => true,
}
}
#[cfg(windows)]
fn is_process_alive(pid: u32) -> bool {
use windows::Win32::Foundation::{CloseHandle, WAIT_TIMEOUT};
use windows::Win32::System::Threading::{
OpenProcess, PROCESS_SYNCHRONIZE, WaitForSingleObject,
};
// SAFETY: OpenProcess returns Err on absence/permission failure;
// PROCESS_SYNCHRONIZE is the minimum right needed for WaitForSingleObject.
let Ok(handle) = (unsafe { OpenProcess(PROCESS_SYNCHRONIZE, false, pid) }) else {
return false;
};
// SAFETY: handle is valid; timeout 0 means "poll, don't block."
let wait_result = unsafe { WaitForSingleObject(handle, 0) };
// SAFETY: handle is owned by us; close regardless of wait result.
let _ = unsafe { CloseHandle(handle) };
wait_result == WAIT_TIMEOUT
}
pub struct DreamLock {
path: PathBuf,
}
impl DreamLock {
pub fn new(workspace_dir: &Path) -> Self {
Self {
path: workspace_dir.join(LOCK_FILE_NAME),
}
}
/// Read the last consolidation timestamp (lock file mtime).
/// Returns `None` if the lock file doesn't exist.
pub fn last_consolidated_at(&self) -> io::Result<Option<SystemTime>> {
match fs::metadata(&self.path) {
Ok(meta) => Ok(Some(meta.modified()?)),
Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(None),
Err(e) => Err(e),
}
}
/// Try to acquire the lock for consolidation.
///
/// Returns `Ok(Some(prior))` on success, where `prior` is the previous mtime
/// (`None` if the file didn't exist). Pass to [`Self::rollback`] on failure.
/// Returns `Ok(None)` if held by a live, non-stale process.
///
/// Reclaims stale locks when the holder PID is dead or age exceeds `stale_secs`.
///
/// Note: this is best-effort coordination, not mutual exclusion. The
/// write-then-verify protocol reduces but cannot eliminate races — two
/// processes may rarely both believe they acquired. Callers must tolerate
/// duplicate consolidation (dream is idempotent).
pub fn try_acquire(&self, stale_secs: u64) -> io::Result<Option<Option<SystemTime>>> {
let prior = match fs::metadata(&self.path) {
Ok(meta) => {
let mtime = meta.modified()?;
if let Ok(content) = fs::read_to_string(&self.path)
&& let Ok(pid) = content.trim().parse::<u32>()
{
let age = SystemTime::now()
.duration_since(mtime)
.unwrap_or_default()
.as_secs();
if age < stale_secs && is_process_alive(pid) {
return Ok(None);
}
}
Some(mtime)
}
Err(e) if e.kind() == io::ErrorKind::NotFound => None,
Err(e) => return Err(e),
};
if let Some(parent) = self.path.parent() {
fs::create_dir_all(parent)?;
}
let our_pid = std::process::id();
fs::write(&self.path, our_pid.to_string())?;
// Re-read to verify we won the race
let content = fs::read_to_string(&self.path)?;
if content.trim().parse::<u32>().ok() == Some(our_pid) {
Ok(Some(prior))
} else {
Ok(None)
}
}
/// Restore lock state after a failed dream.
/// If `prior` is `None` (no prior file), deletes the lock file.
pub fn rollback(&self, prior: Option<SystemTime>) -> io::Result<()> {
match prior {
None => {
if let Err(e) = fs::remove_file(&self.path)
&& e.kind() != io::ErrorKind::NotFound
{
return Err(e);
}
Ok(())
}
Some(mtime) => {
// Clear the PID body so our alive PID doesn't block future reclaimers.
fs::write(&self.path, "")?;
let file = fs::File::options().write(true).open(&self.path)?;
file.set_times(fs::FileTimes::new().set_modified(mtime))
}
}
}
/// Stamp the lock file with the current time to record a consolidation.
pub fn record_consolidation(&self) -> io::Result<()> {
if let Some(parent) = self.path.parent() {
fs::create_dir_all(parent)?;
}
fs::write(&self.path, std::process::id().to_string())
}
}
/// Count session files modified after `since`, excluding the current session.
///
/// Returns sorted file stems of matching `.md` files in `sessions_dir`.
pub fn sessions_since(
sessions_dir: &Path,
since: SystemTime,
exclude_sid8: Option<&str>,
) -> io::Result<Vec<String>> {
let entries = match fs::read_dir(sessions_dir) {
Ok(entries) => entries,
Err(e) if e.kind() == io::ErrorKind::NotFound => return Ok(Vec::new()),
Err(e) => return Err(e),
};
let mut result = Vec::new();
for entry in entries {
let entry = entry?;
let path = entry.path();
if path.extension().and_then(|e| e.to_str()) != Some("md") {
continue;
}
if let Some(exclude) = exclude_sid8
&& path
.file_stem()
.and_then(|s| s.to_str())
.is_some_and(|stem| stem.ends_with(exclude))
{
continue;
}
if entry.metadata()?.modified()? > since
&& let Some(stem) = path.file_stem().and_then(|s| s.to_str())
{
result.push(stem.to_owned());
}
}
result.sort();
Ok(result)
}
#[cfg(test)]
mod tests {
use super::*;
use filetime::FileTime;
use std::time::Duration;
use tempfile::TempDir;
// --- DreamLock tests ---
#[test]
fn no_file_means_no_prior_consolidation() {
let dir = TempDir::new().unwrap();
let lock = DreamLock::new(dir.path());
assert!(lock.last_consolidated_at().unwrap().is_none());
}
#[test]
fn acquire_on_empty_dir_writes_pid() {
let dir = TempDir::new().unwrap();
let lock = DreamLock::new(dir.path());
let prior = lock.try_acquire(300).unwrap().expect("should acquire");
assert!(prior.is_none(), "no prior file existed");
let content = fs::read_to_string(&lock.path).unwrap();
assert_eq!(content, std::process::id().to_string());
assert!(lock.last_consolidated_at().unwrap().is_some());
}
#[test]
fn rollback_none_deletes_file() {
let dir = TempDir::new().unwrap();
let lock = DreamLock::new(dir.path());
let prior = lock.try_acquire(300).unwrap().unwrap();
assert!(lock.path.exists());
lock.rollback(prior).unwrap();
assert!(!lock.path.exists());
assert!(lock.last_consolidated_at().unwrap().is_none());
}
#[test]
fn rollback_restores_prior_mtime() {
let dir = TempDir::new().unwrap();
let lock = DreamLock::new(dir.path());
let old_time = SystemTime::now() - Duration::from_secs(7200);
fs::write(&lock.path, "4000000000").unwrap(); // dead PID
filetime::set_file_mtime(&lock.path, FileTime::from_system_time(old_time)).unwrap();
let prior = lock
.try_acquire(300)
.unwrap()
.expect("should reclaim dead PID");
let prior_mtime = prior.expect("prior file existed");
// mtime after acquire is fresh (from fs::write)
let fresh = lock.last_consolidated_at().unwrap().unwrap();
let fresh_age = SystemTime::now().duration_since(fresh).unwrap_or_default();
assert!(fresh_age.as_secs() < 5);
// Rollback restores old mtime
lock.rollback(Some(prior_mtime)).unwrap();
let restored = lock.last_consolidated_at().unwrap().unwrap();
let drift = restored
.duration_since(old_time)
.or_else(|_| old_time.duration_since(restored))
.unwrap();
assert!(drift.as_secs() < 2, "mtime should be restored");
}
#[test]
fn dead_pid_is_reclaimed() {
let dir = TempDir::new().unwrap();
let lock = DreamLock::new(dir.path());
fs::write(&lock.path, "4000000000").unwrap();
assert!(
lock.try_acquire(300).unwrap().is_some(),
"dead PID should be reclaimable"
);
let content = fs::read_to_string(&lock.path).unwrap();
assert_eq!(content, std::process::id().to_string());
}
#[test]
fn live_pid_blocks_acquisition() {
let dir = TempDir::new().unwrap();
let lock = DreamLock::new(dir.path());
assert!(lock.try_acquire(300).unwrap().is_some(), "first acquire");
assert!(
lock.try_acquire(300).unwrap().is_none(),
"second acquire should be blocked by live PID"
);
}
#[test]
fn stale_age_allows_reclaim_even_if_alive() {
let dir = TempDir::new().unwrap();
let lock = DreamLock::new(dir.path());
fs::write(&lock.path, std::process::id().to_string()).unwrap();
let old = SystemTime::now() - Duration::from_secs(600);
filetime::set_file_mtime(&lock.path, FileTime::from_system_time(old)).unwrap();
// stale_secs=300, age=600 → stale, should reclaim
assert!(
lock.try_acquire(300).unwrap().is_some(),
"stale lock should be reclaimable"
);
}
#[test]
fn record_consolidation_creates_file() {
let dir = TempDir::new().unwrap();
let lock = DreamLock::new(dir.path());
lock.record_consolidation().unwrap();
assert!(lock.path.exists());
let age = SystemTime::now()
.duration_since(lock.last_consolidated_at().unwrap().unwrap())
.unwrap_or_default();
assert!(age.as_secs() < 5);
}
#[test]
fn record_consolidation_updates_mtime() {
let dir = TempDir::new().unwrap();
let lock = DreamLock::new(dir.path());
fs::write(&lock.path, "12345").unwrap();
let old = SystemTime::now() - Duration::from_secs(7200);
filetime::set_file_mtime(&lock.path, FileTime::from_system_time(old)).unwrap();
lock.record_consolidation().unwrap();
let age = SystemTime::now()
.duration_since(lock.last_consolidated_at().unwrap().unwrap())
.unwrap_or_default();
assert!(age.as_secs() < 5, "mtime should be ~now");
}
#[test]
fn full_lifecycle_acquire_consolidate_blocks_reacquire() {
let dir = TempDir::new().unwrap();
let lock = DreamLock::new(dir.path());
let prior = lock.try_acquire(300).unwrap().unwrap();
assert!(prior.is_none());
lock.record_consolidation().unwrap();
assert!(
lock.try_acquire(300).unwrap().is_none(),
"fresh consolidation should block re-acquire"
);
}
#[test]
fn rollback_on_nonexistent_file_is_noop() {
let dir = TempDir::new().unwrap();
let lock = DreamLock::new(dir.path());
lock.rollback(None).unwrap(); // no file to delete, should be fine
}
#[test]
fn corrupted_lock_body_is_reclaimable() {
let dir = TempDir::new().unwrap();
let lock = DreamLock::new(dir.path());
fs::write(&lock.path, "not-a-pid").unwrap();
assert!(
lock.try_acquire(300).unwrap().is_some(),
"unparseable PID should be reclaimable"
);
}
#[test]
fn empty_lock_body_is_reclaimable() {
let dir = TempDir::new().unwrap();
let lock = DreamLock::new(dir.path());
fs::write(&lock.path, "").unwrap();
assert!(
lock.try_acquire(300).unwrap().is_some(),
"empty body should be reclaimable"
);
}
// --- sessions_since tests ---
fn write_session(dir: &Path, name: &str, age_secs: u64) {
fs::create_dir_all(dir).unwrap();
let path = dir.join(format!("{name}.md"));
fs::write(&path, "test").unwrap();
let t = SystemTime::now() - Duration::from_secs(age_secs);
filetime::set_file_mtime(&path, FileTime::from_system_time(t)).unwrap();
}
#[test]
fn filters_by_mtime() {
let dir = TempDir::new().unwrap();
let sessions = dir.path().join("sessions");
let cutoff = SystemTime::now() - Duration::from_secs(3600);
write_session(&sessions, "2026-01-01-proj-aaa11111", 1800); // 30min ago, after cutoff
write_session(&sessions, "2025-12-31-proj-bbb22222", 7200); // 2h ago, before cutoff
let result = sessions_since(&sessions, cutoff, None).unwrap();
assert_eq!(result, vec!["2026-01-01-proj-aaa11111"]);
}
#[test]
fn mtime_at_exact_cutoff_is_excluded() {
let dir = TempDir::new().unwrap();
let sessions = dir.path().join("sessions");
let cutoff = SystemTime::now() - Duration::from_secs(3600);
// Set mtime to the exact cutoff value (not strictly after)
fs::create_dir_all(&sessions).unwrap();
let path = sessions.join("2026-01-01-proj-exact000.md");
fs::write(&path, "test").unwrap();
filetime::set_file_mtime(&path, FileTime::from_system_time(cutoff)).unwrap();
let result = sessions_since(&sessions, cutoff, None).unwrap();
assert!(
result.is_empty(),
"mtime == cutoff should be excluded (strict >)"
);
}
#[test]
fn excludes_current_session() {
let dir = TempDir::new().unwrap();
let sessions = dir.path().join("sessions");
let cutoff = SystemTime::now() - Duration::from_secs(86400);
write_session(&sessions, "2026-01-01-proj-aaa11111", 100);
write_session(&sessions, "2026-01-01-proj-bbb22222", 100);
let result = sessions_since(&sessions, cutoff, Some("bbb22222")).unwrap();
assert_eq!(result, vec!["2026-01-01-proj-aaa11111"]);
}
#[test]
fn empty_dir_returns_empty_vec() {
let dir = TempDir::new().unwrap();
let sessions = dir.path().join("sessions");
fs::create_dir_all(&sessions).unwrap();
let result = sessions_since(&sessions, SystemTime::UNIX_EPOCH, None).unwrap();
assert!(result.is_empty());
}
#[test]
fn nonexistent_dir_returns_empty_vec() {
let dir = TempDir::new().unwrap();
let sessions = dir.path().join("nonexistent");
let result = sessions_since(&sessions, SystemTime::UNIX_EPOCH, None).unwrap();
assert!(result.is_empty());
}
#[test]
fn ignores_non_md_files() {
let dir = TempDir::new().unwrap();
let sessions = dir.path().join("sessions");
fs::create_dir_all(&sessions).unwrap();
write_session(&sessions, "2026-01-01-proj-aaa11111", 0);
fs::write(sessions.join("notes.txt"), "not a session").unwrap();
fs::write(sessions.join("data.json"), "{}").unwrap();
let result = sessions_since(&sessions, SystemTime::UNIX_EPOCH, None).unwrap();
assert_eq!(result, vec!["2026-01-01-proj-aaa11111"]);
}
#[test]
fn returns_sorted_stems() {
let dir = TempDir::new().unwrap();
let sessions = dir.path().join("sessions");
write_session(&sessions, "zzz-session", 0);
write_session(&sessions, "aaa-session", 0);
write_session(&sessions, "mmm-session", 0);
let result = sessions_since(&sessions, SystemTime::UNIX_EPOCH, None).unwrap();
assert_eq!(result, vec!["aaa-session", "mmm-session", "zzz-session"]);
}
}

283
crates/codegen/xai-grok-memory/src/embedding.rs Normal file
View file

@ -0,0 +1,283 @@
//! Embedding provider abstraction for memory vector search.
//!
//! Defines the `EmbeddingProvider` trait and an API-based implementation
//! that calls an OpenAI-compatible embeddings API endpoint.
//!
//! Embeddings are cached in the sqlite-vec `chunks_vec` table — the vec0
//! virtual table IS the cache. No separate cache needed.
use async_trait::async_trait;
/// Maximum retry attempts for transient API errors (429, 5xx).
const MAX_RETRIES: usize = 3;
/// Initial backoff delay in milliseconds (doubles on each retry: 1s, 2s, 4s).
const INITIAL_BACKOFF_MS: u64 = 1000;
/// Trait for generating text embeddings.
///
/// Implementations must be `Send + Sync` so they can be used in `Send`
/// futures (e.g., inside `tokio::spawn`). The `embed_batch` method is
/// async to support API-based providers.
#[async_trait]
pub trait EmbeddingProvider: Send + Sync {
/// Embed a batch of texts, returning one vector per input text.
async fn embed_batch(
&self,
texts: &[&str],
) -> Result<Vec<Vec<f32>>, Box<dyn std::error::Error>>;
/// The model name used for embeddings.
fn model_name(&self) -> &str;
/// The dimensionality of the embedding vectors.
fn dimensions(&self) -> usize;
}
/// API-based embedding provider using an OpenAI-compatible embeddings endpoint.
pub struct ApiEmbeddingProvider {
api_base: String,
model: String,
dimensions: usize,
client: reqwest_middleware::ClientWithMiddleware,
max_batch_size: usize,
}
impl ApiEmbeddingProvider {
pub fn new(
api_base: String,
model: String,
dimensions: usize,
client: reqwest_middleware::ClientWithMiddleware,
) -> Self {
Self {
api_base,
model,
dimensions,
client,
max_batch_size: 32,
}
}
pub fn from_config(
config: &xai_grok_config_types::MemoryEmbeddingConfig,
api_base: String,
client: reqwest_middleware::ClientWithMiddleware,
) -> Option<Self> {
let model = config.model.clone().filter(|m| !m.is_empty())?;
Some(Self::new(api_base, model, config.dimensions, client))
}
pub fn from_session(
config: &xai_grok_config_types::MemoryEmbeddingConfig,
proxy_base_url: String,
auth_key: String,
) -> Option<Self> {
let client = build_static_middleware_client(Some(auth_key));
Self::from_config(config, proxy_base_url, client)
}
}
pub(super) fn build_middleware_client(
credentials: std::sync::Arc<dyn xai_grok_auth::AuthCredentialProvider>,
) -> reqwest_middleware::ClientWithMiddleware {
xai_grok_http::with_auth_retry(xai_grok_http::shared_client(), credentials)
}
fn build_static_middleware_client(
api_key: Option<String>,
) -> reqwest_middleware::ClientWithMiddleware {
let provider: std::sync::Arc<dyn xai_grok_auth::AuthCredentialProvider> = std::sync::Arc::new(
xai_grok_auth::StaticAuthCredentialProvider::new(Box::new(NoopHttpAuth), api_key),
);
build_middleware_client(provider)
}
struct NoopHttpAuth;
impl xai_grok_auth::HttpAuth for NoopHttpAuth {
fn apply(&self, builder: reqwest::RequestBuilder, _base_url: &str) -> reqwest::RequestBuilder {
builder
}
}
#[async_trait]
impl EmbeddingProvider for ApiEmbeddingProvider {
#[tracing::instrument(name = "memory.embed_batch", skip_all, fields(batch_size = texts.len()))]
async fn embed_batch(
&self,
texts: &[&str],
) -> Result<Vec<Vec<f32>>, Box<dyn std::error::Error>> {
if texts.is_empty() {
return Ok(vec![]);
}
let mut all_embeddings = Vec::with_capacity(texts.len());
// Process in batches to respect API payload limits
for batch in texts.chunks(self.max_batch_size) {
let input: Vec<&str> = batch.to_vec();
let body_json = serde_json::json!({
"model": self.model,
"input": input,
"dimensions": self.dimensions,
});
// Retry with exponential backoff on transient errors (429, 5xx)
let mut last_err = String::new();
let mut success = false;
for attempt in 0..MAX_RETRIES {
if attempt > 0 {
let delay = INITIAL_BACKOFF_MS * 2u64.pow(attempt as u32 - 1);
tracing::warn!(
attempt,
delay_ms = delay,
"retrying embedding API call after transient error"
);
tokio::time::sleep(std::time::Duration::from_millis(delay)).await;
}
let request = xai_grok_http::shared_client()
.post(format!("{}/embeddings", self.api_base))
.json(&body_json)
.header("X-XAI-Token-Auth", "xai-grok-cli")
.header("x-grok-client-version", xai_grok_version::VERSION);
let req = match request.build() {
Ok(r) => r,
Err(e) => {
return Err(format!("failed to build embedding request: {e}").into());
}
};
let response = match self.client.execute(req).await {
Ok(r) => r,
Err(e) => {
last_err = format!("request failed: {e}");
continue;
}
};
let status = response.status();
if status.is_success() {
let body: serde_json::Value = response.json().await?;
let data = body
.get("data")
.and_then(|d| d.as_array())
.ok_or("embedding response missing 'data' array")?;
for item in data {
let embedding: Vec<f32> = item
.get("embedding")
.and_then(|e| e.as_array())
.ok_or("embedding item missing 'embedding' array")?
.iter()
.filter_map(|v| v.as_f64().map(|f| f as f32))
.collect();
all_embeddings.push(embedding);
}
success = true;
break;
}
// Retry on 429 (rate limit) or 5xx (server error)
if status == reqwest::StatusCode::TOO_MANY_REQUESTS || status.is_server_error() {
last_err = format!(
"HTTP {status}: {}",
response.text().await.unwrap_or_default()
);
continue;
}
// Non-retryable error (4xx other than 429)
let body = response.text().await.unwrap_or_default();
return Err(format!("embedding API error {status}: {body}").into());
}
if !success {
return Err(format!(
"embedding API failed after {MAX_RETRIES} attempts: {last_err}"
)
.into());
}
}
Ok(all_embeddings)
}
fn model_name(&self) -> &str {
&self.model
}
fn dimensions(&self) -> usize {
self.dimensions
}
}
/// A mock embedding provider for testing that returns deterministic vectors.
/// Uses blake3 hash of text → float values for reproducible results.
#[cfg(test)]
pub struct MockEmbeddingProvider {
pub dimensions: usize,
}
#[cfg(test)]
#[async_trait]
impl EmbeddingProvider for MockEmbeddingProvider {
async fn embed_batch(
&self,
texts: &[&str],
) -> Result<Vec<Vec<f32>>, Box<dyn std::error::Error>> {
Ok(texts
.iter()
.map(|text| {
let hash = blake3::hash(text.as_bytes());
let bytes = hash.as_bytes();
(0..self.dimensions)
.map(|i| bytes[i % 32] as f32 / 255.0)
.collect()
})
.collect())
}
fn model_name(&self) -> &str {
"mock-embedding"
}
fn dimensions(&self) -> usize {
self.dimensions
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_mock_embedding_deterministic() {
let provider = MockEmbeddingProvider { dimensions: 4 };
let r1 = provider.embed_batch(&["hello"]).await.unwrap();
let r2 = provider.embed_batch(&["hello"]).await.unwrap();
assert_eq!(r1, r2);
}
#[tokio::test]
async fn test_mock_embedding_different_texts() {
let provider = MockEmbeddingProvider { dimensions: 4 };
let results = provider.embed_batch(&["hello", "world"]).await.unwrap();
assert_eq!(results.len(), 2);
assert_ne!(results[0], results[1]);
}
#[tokio::test]
async fn test_mock_embedding_empty_input() {
let provider = MockEmbeddingProvider { dimensions: 4 };
let results = provider.embed_batch(&[]).await.unwrap();
assert!(results.is_empty());
}
#[tokio::test]
async fn test_mock_embedding_correct_dimensions() {
let provider = MockEmbeddingProvider { dimensions: 128 };
let results = provider.embed_batch(&["test"]).await.unwrap();
assert_eq!(results[0].len(), 128);
}
}

1286
crates/codegen/xai-grok-memory/src/index.rs Normal file
View file

File diff suppressed because it is too large Load diff

109
crates/codegen/xai-grok-memory/src/lib.rs Normal file
View file

@ -0,0 +1,109 @@
//! Memory system for cross-session knowledge persistence.
//!
//! This crate provides a markdown-based memory storage layer that allows
//! Grok to persist important information across sessions. Memory files are
//! stored under `~/.grok/memory/` with workspace-scoped subdirectories
//! keyed by a blake3 hash of the workspace path.
//!
//! ## Data Layout
//!
//! ```text
//! ~/.grok/memory/
//! ├── MEMORY.md # Global curated knowledge
//! └── {workspace_hash}/ # Per-workspace (blake3(cwd)[..16])
//! ├── MEMORY.md # Project-level curated knowledge
//! └── sessions/
//! └── YYYY-MM-DD-{slug}-{sid8}.md # Session logs
//! ```
//!
//! ## Feature Flag
//!
//! Memory is gated behind `--experimental-memory` CLI flag or
//! `GROK_MEMORY=1` environment variable. When disabled, this crate
//! is not initialized by the host.
pub mod archive;
pub mod backend;
pub mod chunker;
pub mod dream;
pub mod dream_lock;
pub mod embedding;
pub mod index;
pub mod mmr;
pub mod query_expansion;
pub mod schema;
pub mod search;
pub mod storage;
pub mod text_utils;
pub mod watcher;
pub use backend::{MemoryBackendImpl, MemoryBackendParams};
pub use index::{MemoryIndex, init_sqlite_vec};
pub use storage::{MemoryScope, MemoryStorage};
/// Embed all chunks that don't have embeddings yet.
///
/// Queries the index for unembedded chunks, batches them through the
/// embedding provider, and upserts the results. Logs progress.
///
/// This is the async glue between the sync `MemoryIndex` and the async
/// `EmbeddingProvider`. Call after reindex, flush writes, or session-end writes.
pub async fn embed_missing_chunks(
index: &MemoryIndex,
provider: &dyn embedding::EmbeddingProvider,
) -> usize {
let chunks = match index.chunks_without_embeddings() {
Ok(c) if c.is_empty() => return 0,
Ok(c) => c,
Err(e) => {
tracing::warn!(
target: xai_grok_telemetry::memory_log::TARGET,
error = %e,
"failed to query chunks without embeddings"
);
return 0;
}
};
let total = chunks.len();
let mut embedded = 0;
// Batch in groups of 32 (provider's typical max batch size)
for batch in chunks.chunks(32) {
let texts: Vec<&str> = batch.iter().map(|(_, text)| text.as_str()).collect();
match provider.embed_batch(&texts).await {
Ok(embeddings) => {
for ((chunk_id, _), embedding) in batch.iter().zip(embeddings.iter()) {
if let Err(e) = index.upsert_embedding(chunk_id, embedding) {
tracing::warn!(
target: xai_grok_telemetry::memory_log::TARGET,
chunk_id,
error = %e,
"failed to upsert embedding"
);
} else {
embedded += 1;
}
}
}
Err(e) => {
tracing::warn!(
target: xai_grok_telemetry::memory_log::TARGET,
error = %e,
batch_size = texts.len(),
"embedding batch failed, skipping"
);
}
}
}
if embedded > 0 {
tracing::info!(
target: xai_grok_telemetry::memory_log::TARGET,
embedded,
total,
"embedded missing chunks"
);
}
embedded
}

348
crates/codegen/xai-grok-memory/src/mmr.rs Normal file
View file

@ -0,0 +1,348 @@
//! Maximal Marginal Relevance (MMR) diversity re-ranking.
//!
//! Without MMR, if a user has multiple memory chunks about the same topic,
//! the top results are nearly identical. MMR penalizes redundancy by
//! greedily selecting results that balance relevance with diversity.
//!
//! **Formula:**
//! ```text
//! MMR(d) = λ × relevance(d) - (1-λ) × max_similarity(d, selected)
//! ```
//!
//! Uses Jaccard similarity on tokenized snippets (no embeddings needed).
//! O(n²) but n is tiny (typically 6–18 candidates after hybrid scoring).
use std::collections::HashSet;
use super::search::SearchResult;
use xai_grok_config_types::MmrConfig;
/// Tokenize text into a set of alphanumeric words for Jaccard comparison.
///
/// Expects **pre-lowered** input — callers should lowercase snippets before
/// calling this. Uses the same splitting strategy as `query_expansion`
/// (split on non-alphanumeric except underscore) for consistency, but without
/// stop word removal — we want full token overlap for similarity measurement.
fn tokenize(text: &str) -> HashSet<&str> {
text.split(|c: char| !c.is_alphanumeric() && c != '_')
.filter(|w| !w.is_empty())
.collect()
}
/// Jaccard similarity: |A ∩ B| / |A ∪ B|.
fn jaccard_similarity(a: &HashSet<&str>, b: &HashSet<&str>) -> f64 {
if a.is_empty() && b.is_empty() {
return 1.0;
}
if a.is_empty() || b.is_empty() {
return 0.0;
}
let intersection = a.intersection(b).count();
let union = a.len() + b.len() - intersection;
if union == 0 {
0.0
} else {
intersection as f64 / union as f64
}
}
/// Re-rank results using Maximal Marginal Relevance.
///
/// Reorders `results` in-place to balance relevance with diversity.
/// No-op when `config.enabled` is false, `lambda` is 1.0, or there
/// are fewer than 2 results.
///
/// `relevance` is the per-result unclamped ranking score, aligned
/// index-for-index with `results` on entry. It is passed separately rather than
/// read from the clamped `SearchResult.score`, which would saturate top chunks
/// to 1.0 and lose the access-frequency boost tiebreak.
pub fn mmr_rerank(results: &mut Vec<SearchResult>, relevance: &[f64], config: &MmrConfig) {
if !config.enabled || results.len() <= 1 {
return;
}
if config.lambda == 1.0 {
return;
}
assert_eq!(
relevance.len(),
results.len(),
"relevance must be aligned with results"
);
// Lowercase snippets once, then tokenize. This ensures "Rust" and "rust"
// are treated as the same token — casing varies across markdown sources.
let lowered: Vec<String> = results.iter().map(|r| r.snippet.to_lowercase()).collect();
let token_cache: Vec<HashSet<&str>> = lowered.iter().map(|s| tokenize(s)).collect();
let max_score = relevance.iter().copied().fold(f64::NEG_INFINITY, f64::max);
let min_score = relevance.iter().copied().fold(f64::INFINITY, f64::min);
let range = (max_score - min_score).max(f64::EPSILON);
let lambda = config.lambda;
let mut selected: Vec<usize> = Vec::with_capacity(results.len());
let mut remaining: Vec<usize> = (0..results.len()).collect();
while !remaining.is_empty() {
let mut best_pos = 0;
let mut best_mmr = f64::NEG_INFINITY;
for (pos, &candidate) in remaining.iter().enumerate() {
let normalized = (relevance[candidate] - min_score) / range;
let max_sim = selected
.iter()
.map(|&sel| jaccard_similarity(&token_cache[candidate], &token_cache[sel]))
.fold(0.0_f64, f64::max);
let mmr_score = lambda * normalized - (1.0 - lambda) * max_sim;
if mmr_score > best_mmr
|| (mmr_score == best_mmr && relevance[candidate] > relevance[remaining[best_pos]])
{
best_mmr = mmr_score;
best_pos = pos;
}
}
selected.push(remaining.remove(best_pos));
}
let reordered: Vec<SearchResult> = selected
.into_iter()
.map(|i| std::mem::replace(&mut results[i], placeholder_result()))
.collect();
*results = reordered;
// `results` is now reordered, so the caller's `relevance` slice is stale
// and must not be read again.
}
/// Placeholder to enable moving results out of the vec without Clone.
fn placeholder_result() -> SearchResult {
SearchResult {
chunk_id: String::new(),
path: String::new(),
start_line: 0,
end_line: 0,
score: 0.0,
snippet: String::new(),
source: String::new(),
created_at: 0,
}
}
#[cfg(test)]
mod tests {
use super::*;
fn make_result(id: &str, snippet: &str, score: f64) -> SearchResult {
SearchResult {
chunk_id: id.to_string(),
path: format!("{id}.md"),
start_line: 0,
end_line: 1,
score,
snippet: snippet.to_string(),
source: "workspace".to_string(),
created_at: 1_700_000_000,
}
}
fn enabled_config(lambda: f64) -> MmrConfig {
MmrConfig {
enabled: true,
lambda,
}
}
/// Test helper: re-rank using each result's own `score` as its relevance
/// (mirrors the pre-split behavior the existing assertions were written for).
fn rerank(results: &mut Vec<SearchResult>, config: &MmrConfig) {
let relevance: Vec<f64> = results.iter().map(|r| r.score).collect();
mmr_rerank(results, &relevance, config);
}
#[test]
fn test_disabled_is_noop() {
let mut results = vec![
make_result("a", "rust async", 1.0),
make_result("b", "rust async patterns", 0.9),
];
let original_order: Vec<String> = results.iter().map(|r| r.chunk_id.clone()).collect();
rerank(&mut results, &MmrConfig::default());
let after: Vec<String> = results.iter().map(|r| r.chunk_id.clone()).collect();
assert_eq!(original_order, after);
}
#[test]
fn test_lambda_one_is_noop() {
let mut results = vec![
make_result("a", "rust async", 1.0),
make_result("b", "python sync", 0.5),
];
rerank(&mut results, &enabled_config(1.0));
assert_eq!(results[0].chunk_id, "a");
assert_eq!(results[1].chunk_id, "b");
}
#[test]
fn test_single_result_is_noop() {
let mut results = vec![make_result("a", "rust async", 1.0)];
rerank(&mut results, &enabled_config(0.7));
assert_eq!(results.len(), 1);
assert_eq!(results[0].chunk_id, "a");
}
/// Regression guard: MMR must rank on `relevance`, not the clamped
/// `SearchResult.score`. Both results tie at `score == 1.0`; the
/// higher-relevance result is placed SECOND so a buggy `.score` read would
/// keep input order and land "low" first.
#[test]
fn test_mmr_ranks_on_relevance_not_clamped_score() {
let mut results = vec![
make_result("low", "alpha topic one", 1.0),
make_result("high", "beta subject two", 1.0),
];
let relevance = [1.0, 1.25];
mmr_rerank(&mut results, &relevance, &enabled_config(0.7));
assert_eq!(
results[0].chunk_id, "high",
"MMR must order by unclamped relevance, not the clamped .score",
);
assert_eq!(results[1].chunk_id, "low");
}
#[test]
fn test_diverse_results_promoted() {
// Three results: two very similar (rust async), one different (python web)
// With MMR, the diverse result should be promoted over the redundant one
let mut results = vec![
make_result("a", "rust async programming patterns", 1.0),
make_result("b", "rust async programming tutorial", 0.95),
make_result("c", "python web framework flask", 0.9),
];
rerank(&mut results, &enabled_config(0.5));
// First should still be "a" (highest relevance)
assert_eq!(results[0].chunk_id, "a");
// "c" (diverse) should be promoted above "b" (redundant with "a")
assert_eq!(
results[1].chunk_id, "c",
"diverse result should be promoted over redundant one"
);
assert_eq!(results[2].chunk_id, "b");
}
#[test]
fn test_identical_snippets_heavily_penalized() {
let mut results = vec![
make_result("a", "exact same content here", 1.0),
make_result("b", "exact same content here", 0.99),
make_result("c", "completely different topic", 0.5),
];
rerank(&mut results, &enabled_config(0.5));
assert_eq!(results[0].chunk_id, "a");
// "c" should beat "b" because "b" is identical to "a"
assert_eq!(
results[1].chunk_id, "c",
"different result should beat identical duplicate"
);
}
#[test]
fn test_case_insensitive_similarity() {
// "Rust Async" and "rust async" should be treated as identical
// (both lowercased before tokenization). Without lowercasing,
// these would only have 0.5 Jaccard similarity.
let mut results = vec![
make_result("a", "Rust Async Programming", 1.0),
make_result("b", "rust async programming", 0.95),
make_result("c", "Python Web Framework", 0.9),
];
rerank(&mut results, &enabled_config(0.5));
assert_eq!(results[0].chunk_id, "a");
// "c" (diverse) should beat "b" (same content, different casing)
assert_eq!(
results[1].chunk_id, "c",
"case-only difference should be detected as redundant"
);
}
#[test]
fn test_preserves_result_count() {
let mut results = vec![
make_result("a", "one", 1.0),
make_result("b", "two", 0.9),
make_result("c", "three", 0.8),
make_result("d", "four", 0.7),
];
rerank(&mut results, &enabled_config(0.7));
assert_eq!(results.len(), 4);
}
#[test]
fn test_scores_and_snippets_preserved() {
let mut results = vec![
make_result("a", "rust programming", 1.0),
make_result("b", "python scripting", 0.5),
];
rerank(&mut results, &enabled_config(0.7));
// All fields should be intact after re-ranking
for r in &results {
assert!(!r.chunk_id.is_empty());
assert!(!r.snippet.is_empty());
assert!(r.score > 0.0);
}
}
// -----------------------------------------------------------------------
// Jaccard similarity unit tests
// -----------------------------------------------------------------------
#[test]
fn test_jaccard_identical() {
let a: HashSet<&str> = ["rust", "async"].into();
let b: HashSet<&str> = ["rust", "async"].into();
assert!((jaccard_similarity(&a, &b) - 1.0).abs() < f64::EPSILON);
}
#[test]
fn test_jaccard_disjoint() {
let a: HashSet<&str> = ["rust", "async"].into();
let b: HashSet<&str> = ["python", "web"].into();
assert!((jaccard_similarity(&a, &b)).abs() < f64::EPSILON);
}
#[test]
fn test_jaccard_partial_overlap() {
let a: HashSet<&str> = ["rust", "async", "programming"].into();
let b: HashSet<&str> = ["rust", "web", "programming"].into();
// intersection = {rust, programming} = 2, union = {rust, async, programming, web} = 4
assert!((jaccard_similarity(&a, &b) - 0.5).abs() < f64::EPSILON);
}
#[test]
fn test_jaccard_both_empty() {
let a: HashSet<&str> = HashSet::new();
let b: HashSet<&str> = HashSet::new();
assert!((jaccard_similarity(&a, &b) - 1.0).abs() < f64::EPSILON);
}
#[test]
fn test_jaccard_one_empty() {
let a: HashSet<&str> = ["rust"].into();
let b: HashSet<&str> = HashSet::new();
assert!((jaccard_similarity(&a, &b)).abs() < f64::EPSILON);
}
#[test]
fn test_tokenize_splits_on_punctuation() {
let tokens = tokenize("hello, world! rust_code");
assert!(tokens.contains("hello"));
assert!(tokens.contains("world"));
assert!(tokens.contains("rust_code"));
assert!(!tokens.contains(","));
}
}

279
crates/codegen/xai-grok-memory/src/query_expansion.rs Normal file
View file

@ -0,0 +1,279 @@
//! Query expansion for FTS-only search mode.
//!
//! When users ask conversational queries like *"that thing we discussed about the API"*,
//! FTS5 matches every word equally — articles, pronouns, and vague references dilute
//! precision. This module extracts meaningful keywords by removing stop words.
//!
//! The pipeline:
//! ```text
//! query → lowercase → split on non-alphanumeric → remove stop words → dedup → keywords
//! ```
//!
//! When all words are stop words (e.g. "what is that?"), returns an empty vec.
//! The caller (hybrid search) falls back to the vector path in that case.
use std::collections::HashSet;
use std::sync::LazyLock;
static STOP_WORDS: LazyLock<HashSet<&'static str>> = LazyLock::new(|| {
[
// Articles & determiners
"a",
"an",
"the",
"this",
"that",
"these",
"those",
// Pronouns
"i",
"me",
"my",
"we",
"our",
"you",
"your",
"he",
"she",
"it",
"they",
"him",
"her",
"its",
"them",
"us",
// Common verbs
"is",
"are",
"was",
"were",
"be",
"been",
"being",
"have",
"has",
"had",
"do",
"does",
"did",
"will",
"would",
"could",
"should",
"can",
"may",
"might",
// Prepositions
"in",
"on",
"at",
"to",
"for",
"of",
"with",
"by",
"from",
"about",
"into",
"through",
"during",
"before",
"after",
"above",
"below",
// Conjunctions
"and",
"or",
"but",
"if",
"then",
"because",
"as",
"while",
"when",
"where",
"what",
"which",
"who",
"how",
"why",
// Vague references
"thing",
"things",
"stuff",
"something",
"anything",
"everything",
"one",
"some",
"any",
"all",
"each",
"every",
"both",
"few",
"more",
// Time references
"yesterday",
"today",
"tomorrow",
"earlier",
"later",
"recently",
"now",
"just",
"already",
"still",
"yet",
// Request words
"please",
"help",
"find",
"show",
"get",
"tell",
"give",
"make",
// Common filler
"not",
"no",
"yes",
"also",
"too",
"very",
"really",
"here",
"there",
"so",
"up",
"out",
"like",
"than",
"other",
"only",
]
.into_iter()
.collect()
});
/// Extract meaningful keywords from a conversational query by removing stop words.
///
/// Returns keywords in order of appearance, deduplicated. Words shorter than
/// 2 characters and pure-numeric tokens are filtered out. The 2-char minimum
/// preserves meaningful short terms like "go", "js", "ui", "db", "ai", "ml"
/// while stop words handle the common 2-letter noise ("is", "it", "do", "we").
///
/// Returns an empty vec when all words are stop words or the query contains
/// no meaningful content — the caller should fall back to vector search.
pub fn extract_keywords(query: &str) -> Vec<String> {
let lowered = query.to_lowercase();
let mut seen = HashSet::new();
lowered
.split(|c: char| !c.is_alphanumeric() && c != '_')
.filter(|w| w.len() >= 2)
.filter(|w| !STOP_WORDS.contains(w))
.filter(|w| !w.chars().all(|c| c.is_numeric()))
.filter(|w| seen.insert(*w))
.map(|w| w.to_string())
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_removes_stop_words() {
let kw = extract_keywords("that thing we discussed about the API");
assert_eq!(kw, vec!["discussed", "api"]);
}
#[test]
fn test_all_stop_words_returns_empty() {
let kw = extract_keywords("what is that?");
assert!(kw.is_empty());
}
#[test]
fn test_preserves_meaningful_words() {
let kw = extract_keywords("rust programming async patterns");
assert_eq!(kw, vec!["rust", "programming", "async", "patterns"]);
}
#[test]
fn test_filters_single_char_words() {
let kw = extract_keywords("I a x language");
// "i" = 1 char, "a" = 1 char, "x" = 1 char → all filtered by length
assert_eq!(kw, vec!["language"]);
}
#[test]
fn test_preserves_short_meaningful_terms() {
// 2-char terms that are meaningful in programming should survive
let kw = extract_keywords("Go and JS patterns");
assert_eq!(kw, vec!["go", "js", "patterns"]);
}
#[test]
fn test_short_stop_words_filtered() {
// 2-char stop words ("is", "it", "do", "we") should still be removed
let kw = extract_keywords("is it ok to do that");
assert_eq!(kw, vec!["ok"]);
}
#[test]
fn test_filters_pure_numbers() {
let kw = extract_keywords("port 8080 and 443 config");
assert_eq!(kw, vec!["port", "config"]);
}
#[test]
fn test_deduplicates() {
let kw = extract_keywords("rust rust rust programming");
assert_eq!(kw, vec!["rust", "programming"]);
}
#[test]
fn test_handles_punctuation() {
let kw = extract_keywords("what's the solution for the bug?");
assert_eq!(kw, vec!["solution", "bug"]);
}
#[test]
fn test_preserves_underscored_identifiers() {
let kw = extract_keywords("the my_function variable");
assert_eq!(kw, vec!["my_function", "variable"]);
}
#[test]
fn test_empty_query() {
assert!(extract_keywords("").is_empty());
}
#[test]
fn test_only_punctuation() {
assert!(extract_keywords("??? !!! ...").is_empty());
}
#[test]
fn test_mixed_case() {
let kw = extract_keywords("Rust Programming ASYNC");
assert_eq!(kw, vec!["rust", "programming", "async"]);
}
#[test]
fn test_real_conversational_queries() {
assert_eq!(
extract_keywords("what was the solution for the authentication bug"),
vec!["solution", "authentication", "bug"]
);
assert_eq!(
extract_keywords("how do I configure the memory system"),
vec!["configure", "memory", "system"]
);
assert_eq!(
extract_keywords("show me that database migration we talked about"),
vec!["database", "migration", "talked"]
);
}
}

98
crates/codegen/xai-grok-memory/src/schema.rs Normal file
View file

@ -0,0 +1,98 @@
//! SQL schema constants for the memory index.
//!
//! The index uses three tables:
//! - `meta` — key-value metadata (embedding dimensions, schema version)
//! - `chunks` — indexed text chunks with blake3 content hashes
//! - `chunks_fts` — contentless FTS5 virtual table for BM25 keyword search
//!
//! When sqlite-vec is available, a fourth table is created:
//! - `chunks_vec` — vec0 virtual table for KNN vector search
/// Schema version. Bump when making breaking schema changes that require
/// dropping and recreating tables.
pub const SCHEMA_VERSION: u32 = 1;
/// Generate the SQL schema for the memory index.
///
/// `dimensions` controls the embedding vector size for `chunks_vec`.
/// If `vec_available` is false, the `chunks_vec` table is not created.
///
/// Connection pragmas (busy_timeout, journal_mode) are applied on the open
/// path (`xai_sqlite_journal::JournalMode::open`) — the journal mode depends
/// on the database's filesystem.
pub fn schema_sql(dimensions: usize, vec_available: bool) -> String {
let mut sql = format!(
r#"
CREATE TABLE IF NOT EXISTS meta (
key TEXT PRIMARY KEY,
value TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS chunks (
rowid INTEGER PRIMARY KEY AUTOINCREMENT,
id TEXT UNIQUE NOT NULL,
path TEXT NOT NULL,
start_line INTEGER NOT NULL,
end_line INTEGER NOT NULL,
text TEXT NOT NULL,
hash TEXT NOT NULL,
source TEXT NOT NULL,
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL,
access_count INTEGER DEFAULT 0,
last_accessed INTEGER
);
CREATE INDEX IF NOT EXISTS idx_chunks_path ON chunks(path);
CREATE INDEX IF NOT EXISTS idx_chunks_hash ON chunks(hash);
CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(text, content='');
INSERT OR IGNORE INTO meta(key, value) VALUES ('reindex_claim', '');
"#
);
if vec_available {
sql.push_str(&format!(
"\nCREATE VIRTUAL TABLE IF NOT EXISTS chunks_vec USING vec0(\n \
chunk_id TEXT PRIMARY KEY,\n \
embedding FLOAT[{dimensions}]\n);\n"
));
}
sql
}
/// SQL to insert or update an embedding dimension record in the meta table.
pub const UPSERT_META_SQL: &str = "INSERT OR REPLACE INTO meta(key, value) VALUES (?1, ?2)";
/// SQL to query a meta value by key.
pub const GET_META_SQL: &str = "SELECT value FROM meta WHERE key = ?1";
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_schema_sql_without_vec() {
let sql = schema_sql(1536, false);
assert!(sql.contains("CREATE TABLE IF NOT EXISTS chunks"));
assert!(sql.contains("CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts"));
assert!(!sql.contains("chunks_vec"));
// Connection pragmas live on the open path, not in the schema batch.
assert!(!sql.contains("PRAGMA"));
}
#[test]
fn test_schema_sql_with_vec() {
let sql = schema_sql(384, true);
assert!(sql.contains("chunks_vec"));
assert!(sql.contains("FLOAT[384]"));
}
#[test]
fn test_schema_sql_different_dimensions() {
let sql = schema_sql(768, true);
assert!(sql.contains("FLOAT[768]"));
}
}

1349
crates/codegen/xai-grok-memory/src/search.rs Normal file
View file

File diff suppressed because it is too large Load diff

1862
crates/codegen/xai-grok-memory/src/storage.rs Normal file
View file

File diff suppressed because it is too large Load diff

54
crates/codegen/xai-grok-memory/src/text_utils.rs Normal file
View file

@ -0,0 +1,54 @@
//! Pure text-classification helpers shared by the memory flush
//! (`session::helpers::memory_flush`) and dream (`session::memory::dream`)
//! response-processing paths.
//!
//! These live here, in the memory subsystem, so `dream` no longer reaches
//! *up* into `session::helpers::memory_flush` for them — which removes the
//! `dream` <-> `memory_flush` module dependency cycle and is a prerequisite
//! for extracting the memory subsystem into its own crate.
/// Check if text contains at least one markdown header (`#` or `##`).
///
/// Used by both flush and dream response processing to ensure the model
/// produced structured output.
pub fn has_markdown_headers(text: &str) -> bool {
text.contains("## ") || text.contains("# ")
}
/// Check if the response matches the NO_REPLY convention.
///
/// Strips all non-alphanumeric characters, lowercases, and checks if the
/// remainder is exactly `"noreply"`. This handles common separator variants:
/// `"no reply"`, `"no_reply"`, `"no-reply"`, `"NO REPLY"`, etc.
pub fn is_no_reply(text: &str) -> bool {
let normalized: String = text
.to_lowercase()
.chars()
.filter(|c| c.is_alphanumeric())
.collect();
normalized == "noreply"
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_is_no_reply() {
assert!(is_no_reply("NO_REPLY"));
assert!(is_no_reply("no reply"));
assert!(is_no_reply("No-Reply"));
assert!(is_no_reply("noreply"));
assert!(!is_no_reply("no reply needed"));
assert!(!is_no_reply("I have things to store"));
}
#[test]
fn test_has_markdown_headers() {
assert!(has_markdown_headers("## Topic"));
assert!(has_markdown_headers("# Title\n\nBody"));
assert!(has_markdown_headers("preamble\n\n## Topic"));
assert!(!has_markdown_headers("plain text without headers"));
assert!(!has_markdown_headers("#hashtag without space"));
}
}

192
crates/codegen/xai-grok-memory/src/watcher.rs Normal file
View file

@ -0,0 +1,192 @@
//! File watcher for detecting external memory edits.
//!
//! Watches `~/.grok/memory/` for `.md` file changes (create, modify, remove)
//! and accumulates the affected paths. The search path checks [`is_dirty`]
//! before each query and syncs the index for all dirty paths:
//! - **created / modified** files are reindexed via `MemoryIndex::reindex_file`
//! - **deleted** files have their stale chunks removed via `MemoryIndex::delete_path`
//!
//! Without the deletion handling, chunks from removed files would remain
//! searchable indefinitely.
//!
//! Uses `arc_swap::ArcSwap` for lock-free dirty path tracking — the notify
//! event handler inserts via `rcu`, the search path takes via atomic swap.
//!
//! [`is_dirty`]: MemoryFileWatcher::is_dirty
use std::collections::HashSet;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use arc_swap::ArcSwap;
use notify::{Event, EventKind, RecommendedWatcher, RecursiveMode, Watcher};
/// Watches the memory directory for `.md` file changes.
///
/// Lock-free design:
/// - **Insert** (notify thread): `dirty_files.rcu(|old| { clone + insert })`
/// - **Take** (search path): `dirty_files.swap(empty)` — single atomic pointer exchange
/// - **Quick check**: `dirty.load(Relaxed)` — single atomic load, no allocation
pub struct MemoryFileWatcher {
dirty_files: Arc<ArcSwap<HashSet<PathBuf>>>,
dirty: Arc<AtomicBool>,
_watcher: RecommendedWatcher,
}
impl MemoryFileWatcher {
/// Start watching the given memory directory for `.md` file changes.
///
/// Returns `None` if the watcher fails to initialize (logged, non-fatal).
pub fn start(memory_dir: &Path) -> Option<Self> {
let dirty_files: Arc<ArcSwap<HashSet<PathBuf>>> =
Arc::new(ArcSwap::new(Arc::new(HashSet::new())));
let dirty = Arc::new(AtomicBool::new(false));
let df = dirty_files.clone();
let d = dirty.clone();
let mut watcher = notify::recommended_watcher(move |res: Result<Event, notify::Error>| {
let Ok(event) = res else { return };
match event.kind {
EventKind::Create(_) | EventKind::Modify(_) | EventKind::Remove(_) => {}
_ => return,
}
for path in &event.paths {
if path.extension().is_some_and(|ext| ext == "md") {
let path = path.clone();
df.rcu(move |old| {
let mut new = (**old).clone();
new.insert(path.clone());
new
});
d.store(true, Ordering::Relaxed);
}
}
})
.map_err(|e| {
tracing::warn!(error = %e, "failed to create memory file watcher");
})
.ok()?;
watcher
.watch(memory_dir, RecursiveMode::Recursive)
.map_err(|e| {
tracing::warn!(
path = %memory_dir.display(),
error = %e,
"failed to watch memory directory"
);
})
.ok()?;
tracing::info!(
path = %memory_dir.display(),
"memory file watcher started"
);
Some(Self {
dirty_files,
dirty,
_watcher: watcher,
})
}
/// Quick check: true if any files have been modified since last take.
pub fn is_dirty(&self) -> bool {
self.dirty.load(Ordering::Relaxed)
}
/// Take all accumulated dirty paths, resetting the dirty state.
/// Returns the paths that changed since the last take.
pub fn take_dirty(&self) -> Vec<PathBuf> {
let old = self.dirty_files.swap(Arc::new(HashSet::new()));
self.dirty.store(false, Ordering::Relaxed);
old.iter().cloned().collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
#[test]
fn test_watcher_starts_on_valid_dir() {
let tmp = TempDir::new().unwrap();
// In CI / containerized environments the OS may deny inotify watches
// (e.g. exhausted fs.inotify.max_user_instances); skip gracefully.
let _watcher = MemoryFileWatcher::start(tmp.path());
}
#[test]
fn test_watcher_initially_clean() {
let tmp = TempDir::new().unwrap();
let Some(watcher) = MemoryFileWatcher::start(tmp.path()) else {
eprintln!("skipping: could not create file watcher (resource limit?)");
return;
};
assert!(!watcher.is_dirty());
assert!(watcher.take_dirty().is_empty());
}
#[test]
fn test_watcher_detects_md_file_creation() {
let tmp = TempDir::new().unwrap();
let Some(watcher) = MemoryFileWatcher::start(tmp.path()) else {
eprintln!("skipping: could not create file watcher (resource limit?)");
return;
};
// Create a .md file — watcher should detect it
std::fs::write(tmp.path().join("test.md"), "hello").unwrap();
// Give the watcher time to process (debounce + OS event delivery)
std::thread::sleep(std::time::Duration::from_millis(500));
assert!(watcher.is_dirty(), "should detect .md creation");
let dirty = watcher.take_dirty();
assert!(!dirty.is_empty(), "should have dirty paths");
assert!(dirty[0].extension().unwrap() == "md");
}
#[test]
fn test_watcher_ignores_non_md_files() {
let tmp = TempDir::new().unwrap();
let Some(watcher) = MemoryFileWatcher::start(tmp.path()) else {
eprintln!("skipping: could not create file watcher (resource limit?)");
return;
};
// Create a non-.md file
std::fs::write(tmp.path().join("test.txt"), "hello").unwrap();
std::fs::write(tmp.path().join("index.sqlite"), "db").unwrap();
std::thread::sleep(std::time::Duration::from_millis(500));
assert!(
!watcher.is_dirty(),
"should not detect non-.md file changes"
);
}
#[test]
fn test_take_dirty_resets_state() {
let tmp = TempDir::new().unwrap();
let Some(watcher) = MemoryFileWatcher::start(tmp.path()) else {
eprintln!("skipping: could not create file watcher (resource limit?)");
return;
};
std::fs::write(tmp.path().join("a.md"), "content").unwrap();
std::thread::sleep(std::time::Duration::from_millis(500));
let first = watcher.take_dirty();
assert!(!first.is_empty());
assert!(!watcher.is_dirty(), "should be clean after take");
assert!(
watcher.take_dirty().is_empty(),
"second take should be empty"
);
}
}