From edb84bdd46b0cbc3d8e4cfe7a96be833c6878d00 Mon Sep 17 00:00:00 2001 From: Philippe Branchu Date: Wed, 3 Jun 2026 04:47:59 +0000 Subject: [PATCH 01/12] feat(memory): persistent default user + session user-tagging foundation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Lays the groundwork for per-user memory by giving every install a stable default-user UUID and tagging every session with an owning user. Sessions are now consistently user-scoped: - `Session::user_id: UserId` (required, not Option) — defaults to the kernel's persistent default user - `Session::parent_session_id: Option` — foundation for future tree-scoped cascade deletion of forked sessions (no producer yet) - `MessageSource` enum + optional `Message::source` — additive type that later PRs (structured extraction filtering) will read; no consumer here - `UserConfig::is_default: bool` — `[[users]]` blocks can attach display name and channel bindings to the persistent default identity Kernel boots the default user once and caches it process-wide: - `bootstrap_default_user` — load-or-generate the UUID from `kv_store[shared, "default_user_uuid"]`, install via `set_default_user_id`, then run a one-shot rewrite of legacy nil-UUID sessions, gated by the `default_user_bootstrap_done` sentinel - `resolve_user_id` (strict, HTTP boundary) — folds the deprecated "test" alias and the nil UUID to the default user with `warn!` logs so reserved-bucket abuse is auditable - `resolve_user_id_internal` (raw mapper) — preserves the pre-fix behaviour for in-process test callers - `AuthManager::new_with_default` — binds the `is_default = true` user (or the first user) to the persistent UUID Storage and migration: - Schema v9 adds `user_id` (NOT NULL, default nil UUID) and `parent_session_id` (nullable) to `sessions`, plus `(agent_id, user_id)` and `parent_session_id` indexes - `MemorySubstrate::rewrite_nil_user_sessions` — atomic transaction wrapping the legacy-bucket UPDATE; the kernel only sets the bootstrap-done sentinel after a clean rewrite, so a failure leaves the retry path intact Single-user installs see no behaviour change: everyone is the default user, one session per agent, same as today. Tests: - `MessageSource` deserialises cleanly from pre-field payloads (JSON + msgpack) and survives full round-trips - `UserConfig::is_default` defaults to `false` for existing configs - `AuthManager::new_with_default` honours `is_default`, falls back to first user, and is a no-op for empty configs - Migration v9 adds the columns and indexes, and a v8-built DB upgrades cleanly to v9 with pre-existing rows preserved - `default_user_id`/`test_user_id` are distinct - `create_session(agent, user)` round-trips through SQLite — both the default and an explicit user - `rewrite_nil_user_sessions` is idempotent, targeted, and atomic - Kernel: default-user UUID persists across kernel restarts; the strict filter folds `"test"` and nil to default while passing other UUIDs through; the internal mapper preserves the raw `"test"` alias Co-Authored-By: Claude Sonnet 4.6 --- Cargo.lock | 1 + crates/openfang-api/src/openai_compat.rs | 1 + crates/openfang-api/src/routes.rs | 2 + crates/openfang-kernel/Cargo.toml | 1 + crates/openfang-kernel/src/auth.rs | 125 +++++- crates/openfang-kernel/src/kernel.rs | 461 +++++++++++++++++++++- crates/openfang-memory/src/migration.rs | 114 +++++- crates/openfang-memory/src/session.rs | 225 ++++++++++- crates/openfang-memory/src/substrate.rs | 153 ++++++- crates/openfang-runtime/src/agent_loop.rs | 24 ++ crates/openfang-runtime/src/compactor.rs | 12 + crates/openfang-types/src/config.rs | 39 ++ crates/openfang-types/src/message.rs | 85 ++++ 13 files changed, 1209 insertions(+), 34 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e880570247..9903cae3b9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4164,6 +4164,7 @@ dependencies = [ "openfang-wire", "rand 0.8.5", "reqwest 0.12.28", + "rusqlite", "rustls", "serde", "serde_json", diff --git a/crates/openfang-api/src/openai_compat.rs b/crates/openfang-api/src/openai_compat.rs index b87d6893f2..29a55c5a34 100644 --- a/crates/openfang-api/src/openai_compat.rs +++ b/crates/openfang-api/src/openai_compat.rs @@ -240,6 +240,7 @@ fn convert_messages(oai_messages: &[OaiMessage]) -> Vec { provider_msg_id: None, role, content, + source: None, }) }) .collect() diff --git a/crates/openfang-api/src/routes.rs b/crates/openfang-api/src/routes.rs index cebb1f599a..7bb7beecf5 100644 --- a/crates/openfang-api/src/routes.rs +++ b/crates/openfang-api/src/routes.rs @@ -319,6 +319,8 @@ pub fn inject_attachments_into_session( _ => openfang_memory::session::Session { id: entry.session_id, agent_id, + user_id: openfang_memory::session::default_user_id(), + parent_session_id: None, messages: Vec::new(), context_window_tokens: 0, label: None, diff --git a/crates/openfang-kernel/Cargo.toml b/crates/openfang-kernel/Cargo.toml index 851cbfa4d4..fcd175b9c2 100644 --- a/crates/openfang-kernel/Cargo.toml +++ b/crates/openfang-kernel/Cargo.toml @@ -44,3 +44,4 @@ libc = "0.2" [dev-dependencies] tokio-test = { workspace = true } tempfile = { workspace = true } +rusqlite = { workspace = true } diff --git a/crates/openfang-kernel/src/auth.rs b/crates/openfang-kernel/src/auth.rs index 6bc4a538c2..0707178e92 100644 --- a/crates/openfang-kernel/src/auth.rs +++ b/crates/openfang-kernel/src/auth.rs @@ -104,14 +104,52 @@ pub struct AuthManager { impl AuthManager { /// Create a new AuthManager from kernel user configuration. + /// + /// Thin wrapper over [`Self::new_with_default`] for callers that do not + /// have a persistent default-user UUID to bind. Each user gets a freshly + /// generated `UserId`. pub fn new(user_configs: &[UserConfig]) -> Self { + Self::new_with_default(user_configs, None) + } + + /// Like [`Self::new`] but accepts an explicit persistent default-user UUID. + /// + /// When `default_user_id` is `Some(uid)`, the config user marked + /// `is_default = true` (or, if none is marked, the first user in the list) + /// is registered under `uid` instead of a freshly-generated UUID. This + /// binds the user's display name and channel bindings to the kernel's + /// persistent default-session bucket, so unattributed traffic is + /// attributed to a real person rather than the nil-UUID anonymous bucket. + pub fn new_with_default(user_configs: &[UserConfig], default_user_id: Option) -> Self { let manager = Self { users: DashMap::new(), channel_index: DashMap::new(), }; - for config in user_configs { - let user_id = UserId::new(); + // Determine which config index claims the default user identity. + // Priority: an explicit `is_default = true` block, else the first + // user. We only need a `Some` here when the caller actually has a + // persistent UUID to bind, otherwise every user gets a fresh ID. + let default_index: Option = if default_user_id.is_some() { + user_configs + .iter() + .position(|c| c.is_default) + .or(if user_configs.is_empty() { + None + } else { + Some(0) + }) + } else { + None + }; + + for (idx, config) in user_configs.iter().enumerate() { + let user_id = if Some(idx) == default_index { + // SAFETY: default_index is only Some when default_user_id is Some. + default_user_id.unwrap() + } else { + UserId::new() + }; let role = UserRole::from_str_role(&config.role); let identity = UserIdentity { id: user_id, @@ -131,6 +169,7 @@ impl AuthManager { user = %config.name, role = %role, bindings = config.channel_bindings.len(), + default = (Some(idx) == default_index), "Registered user" ); } @@ -205,6 +244,7 @@ mod tests { m }, api_key_hash: None, + is_default: false, }, UserConfig { name: "Guest".to_string(), @@ -215,12 +255,14 @@ mod tests { m }, api_key_hash: None, + is_default: false, }, UserConfig { name: "ReadOnly".to_string(), role: "viewer".to_string(), channel_bindings: HashMap::new(), api_key_hash: None, + is_default: false, }, ] } @@ -304,6 +346,85 @@ mod tests { assert_eq!(manager.user_count(), 0); } + #[test] + fn test_new_with_default_binds_is_default_user() { + // When one [[users]] entry sets `is_default = true`, the manager must + // register that user under the persistent UUID (not a fresh one). + let default_uid = UserId::new(); + let configs = vec![ + UserConfig { + name: "Alice".to_string(), + role: "user".to_string(), + channel_bindings: HashMap::new(), + api_key_hash: None, + is_default: false, + }, + UserConfig { + name: "Owner".to_string(), + role: "owner".to_string(), + channel_bindings: HashMap::new(), + api_key_hash: None, + is_default: true, + }, + ]; + + let manager = AuthManager::new_with_default(&configs, Some(default_uid)); + // Owner's name resolves to the persistent UUID. + let owner = manager.get_user(default_uid).expect("owner registered"); + assert_eq!(owner.name, "Owner"); + assert_eq!(owner.role, UserRole::Owner); + } + + #[test] + fn test_new_with_default_falls_back_to_first_user() { + // When no entry sets `is_default`, the first user in the list inherits + // the persistent UUID. This matches the single-user-install default + // where there is exactly one `[[users]]` block. + let default_uid = UserId::new(); + let configs = vec![ + UserConfig { + name: "Primary".to_string(), + role: "owner".to_string(), + channel_bindings: HashMap::new(), + api_key_hash: None, + is_default: false, + }, + UserConfig { + name: "Secondary".to_string(), + role: "user".to_string(), + channel_bindings: HashMap::new(), + api_key_hash: None, + is_default: false, + }, + ]; + + let manager = AuthManager::new_with_default(&configs, Some(default_uid)); + let primary = manager.get_user(default_uid).expect("primary registered"); + assert_eq!(primary.name, "Primary"); + } + + #[test] + fn test_new_with_default_none_keeps_fresh_uuids() { + // Pass `None` (e.g. from a non-kernel test harness): no user is + // bound to the persistent UUID; each entry gets a fresh `UserId`. + let configs = test_configs(); + let manager = AuthManager::new_with_default(&configs, None); + assert_eq!(manager.user_count(), 3); + // None of the users can be looked up under the nil UUID by accident. + assert!(manager.get_user(UserId(uuid::Uuid::nil())).is_none()); + } + + #[test] + fn test_new_with_default_empty_config_is_noop() { + // Empty [[users]] config: the default UUID has nowhere to bind, so + // the manager comes up disabled. This matches a fresh install with + // no user config block. + let default_uid = UserId::new(); + let manager = AuthManager::new_with_default(&[], Some(default_uid)); + assert!(!manager.is_enabled()); + assert_eq!(manager.user_count(), 0); + } + #[test] fn test_role_parsing() { assert_eq!(UserRole::from_str_role("owner"), UserRole::Owner); diff --git a/crates/openfang-kernel/src/kernel.rs b/crates/openfang-kernel/src/kernel.rs index 8f59414c97..f520edeb88 100644 --- a/crates/openfang-kernel/src/kernel.rs +++ b/crates/openfang-kernel/src/kernel.rs @@ -650,6 +650,16 @@ impl OpenFangKernel { .map_err(|e| KernelError::BootFailed(format!("Memory init failed: {e}")))?, ); + // Bootstrap the persistent default user UUID. + // + // Replaces the legacy nil-UUID "anonymous bucket" so unattributed + // sessions get a stable owner. The UUID is generated once on first + // boot and persisted to the shared kv_store; subsequent boots reload + // it. On first boot after upgrading from a pre-v9 schema, any + // existing sessions with `user_id = nil UUID` are rewritten to the + // new default. + Self::bootstrap_default_user(&memory)?; + // Initialize credential resolver (vault → dotenv → env var) let credential_resolver = { let vault_path = config.home_dir.join("vault.enc"); @@ -816,8 +826,18 @@ impl OpenFangKernel { let wasm_sandbox = WasmSandbox::new() .map_err(|e| KernelError::BootFailed(format!("WASM sandbox init failed: {e}")))?; - // Initialize RBAC authentication manager - let auth = AuthManager::new(&config.users); + // Initialize RBAC authentication manager. + // + // The persistent default user (bootstrapped above) is bound to the + // `[[users]]` entry marked `is_default = true`, or to the first + // entry if none is marked. This keeps the default-user identity + // stable across reboots and lets `[[users]]` metadata (display + // name, channel bindings) attach to the persistent bucket rather + // than to a freshly-generated per-config UUID. + let auth = AuthManager::new_with_default( + &config.users, + Some(openfang_memory::session::default_user_id()), + ); if auth.is_enabled() { info!("RBAC enabled with {} users", auth.user_count()); } @@ -1611,9 +1631,11 @@ impl OpenFangKernel { // Create session — use the returned session_id so the registry // and database are in sync (fixes duplicate session bug #651). + // New sessions are tagged with the persistent default user; future + // PRs will surface explicit user selection through this path. let session = self .memory - .create_session(agent_id) + .create_session(agent_id, openfang_memory::session::default_user_id()) .map_err(KernelError::OpenFang)?; let session_id = session.id; @@ -2064,6 +2086,8 @@ impl OpenFangKernel { .unwrap_or_else(|| openfang_memory::session::Session { id: entry.session_id, agent_id, + user_id: openfang_memory::session::default_user_id(), + parent_session_id: None, messages: Vec::new(), context_window_tokens: 0, label: None, @@ -2650,6 +2674,8 @@ impl OpenFangKernel { .unwrap_or_else(|| openfang_memory::session::Session { id: entry.session_id, agent_id, + user_id: openfang_memory::session::default_user_id(), + parent_session_id: None, messages: Vec::new(), context_window_tokens: 0, label: None, @@ -3060,10 +3086,10 @@ impl OpenFangKernel { // Delete the old session let _ = self.memory.delete_session(entry.session_id); - // Create a fresh session + // Create a fresh session bound to the persistent default user. let new_session = self .memory - .create_session(agent_id) + .create_session(agent_id, openfang_memory::session::default_user_id()) .map_err(KernelError::OpenFang)?; // Update registry with new session ID @@ -3092,10 +3118,10 @@ impl OpenFangKernel { // Delete canonical (cross-channel) session let _ = self.memory.delete_canonical_session(agent_id); - // Create a fresh session + // Create a fresh session bound to the persistent default user. let new_session = self .memory - .create_session(agent_id) + .create_session(agent_id, openfang_memory::session::default_user_id()) .map_err(KernelError::OpenFang)?; // Update registry with new session ID @@ -3570,6 +3596,8 @@ impl OpenFangKernel { .unwrap_or_else(|| openfang_memory::session::Session { id: entry.session_id, agent_id, + user_id: openfang_memory::session::default_user_id(), + parent_session_id: None, messages: Vec::new(), context_window_tokens: 0, label: None, @@ -3652,6 +3680,8 @@ impl OpenFangKernel { .unwrap_or_else(|| openfang_memory::session::Session { id: entry.session_id, agent_id, + user_id: openfang_memory::session::default_user_id(), + parent_session_id: None, messages: Vec::new(), context_window_tokens: 0, label: None, @@ -6728,6 +6758,196 @@ impl OpenFangKernel { } } } + + /// Return the persistent default user UUID generated at boot. + pub fn default_user_id(&self) -> openfang_types::agent::UserId { + openfang_memory::session::default_user_id() + } + + /// Public, HTTP-boundary user-ID resolver. + /// + /// This is the strict, security-relevant variant: every reserved-bucket + /// value that an API-key holder could otherwise abuse to attribute writes + /// to is folded back to the persistent default user, with a `warn!` log: + /// + /// - `"default"` → persistent default user + /// - `"test"` → persistent default user (with `warn!` — deprecated alias, + /// external traffic no longer reaches the test bucket) + /// - the nil UUID (`00000000-…`) → persistent default user (with `warn!`) + /// - any other valid UUID → parsed as-is + /// - `None` or unparseable → `None` (callers fall back to the default + /// session selection) + /// + /// Use this for any caller that takes user-controlled input: HTTP route + /// handlers, message dispatchers, channel adapters. Tests and other + /// in-process callers that need to address the well-known test user UUID + /// must call [`Self::resolve_user_id_internal`] instead. + /// + /// # Why this is public in PR 1 + /// + /// No HTTP route consumes this in PR 1 — it is foundation API for PR 2, + /// which adds the `parse_user_id` route helper that delegates here. The + /// strict-filter behaviour is part of the security model and ships best + /// alongside the foundation; downgrading to private would mean + /// re-promoting in PR 2. + pub fn resolve_user_id( + &self, + user_id_str: Option<&str>, + ) -> Option { + let resolved = self.resolve_user_id_internal(user_id_str)?; + let default_user = openfang_memory::session::default_user_id(); + if resolved.0 == openfang_memory::session::TEST_USER_UUID { + warn!( + input = ?user_id_str, + "Rejected reserved \"test\" user alias on external API; mapped to default user" + ); + return Some(default_user); + } + if resolved.0.is_nil() { + warn!( + input = ?user_id_str, + "Rejected nil-UUID user on external API; mapped to default user" + ); + return Some(default_user); + } + Some(resolved) + } + + /// Internal, raw user-ID resolver. + /// + /// Honours `"test"` and the nil UUID literally (no folding), without the + /// security filter applied by [`Self::resolve_user_id`]: + /// + /// - `"test"` → the well-known test user UUID (isolated from + /// default-user memory; see [`openfang_memory::session::TEST_USER_UUID`]) + /// - `"default"` → persistent default user + /// - any valid UUID (including nil) → parsed as-is + /// - `None` or unparseable → `None` + /// + /// Use ONLY from test harnesses or trusted internal code paths. Any + /// caller taking user-controlled input must use the HTTP-boundary + /// [`Self::resolve_user_id`] instead. + /// + /// # Why this is public in PR 1 + /// + /// Same rationale as [`Self::resolve_user_id`]: this is foundation API + /// for the PR 2 route helpers. The internal/strict pair ships together + /// so PR 2 only has to wire the route helper through, without re-opening + /// visibility. + pub fn resolve_user_id_internal( + &self, + user_id_str: Option<&str>, + ) -> Option { + match user_id_str? { + "test" => Some(openfang_types::agent::UserId( + openfang_memory::session::TEST_USER_UUID, + )), + "default" => Some(openfang_memory::session::default_user_id()), + s => s + .parse::() + .ok() + .map(openfang_types::agent::UserId), + } + } + + /// Bootstrap the persistent default user UUID on kernel startup. + /// + /// On first boot: generate a new UUID, persist it under + /// `kv_store[shared_memory_agent_id, "default_user_uuid"]`, and rewrite + /// any legacy nil-UUID sessions to use the new default user. + /// + /// On subsequent boots: load the existing UUID. The rewrite has already + /// happened (sentinel `default_user_bootstrap_done` is set) so it is + /// skipped. + /// + /// The UUID is installed into `openfang_memory::session` via + /// `set_default_user_id`, so every call to `default_user_id()` from + /// anywhere in the codebase returns this persistent value. + fn bootstrap_default_user(memory: &Arc) -> KernelResult<()> { + const KEY_UUID: &str = "default_user_uuid"; + const KEY_DONE: &str = "default_user_bootstrap_done"; + let shared = shared_memory_agent_id(); + + // Load or generate the persistent UUID. + let user_id = match memory + .structured_get(shared, KEY_UUID) + .map_err(|e| KernelError::BootFailed(format!("kv read default_user_uuid: {e}")))? + { + Some(serde_json::Value::String(s)) => match uuid::Uuid::parse_str(&s) { + Ok(u) => { + info!(default_user = %u, "Loaded persistent default user UUID"); + openfang_types::agent::UserId(u) + } + Err(e) => { + warn!( + error = %e, + value = %s, + "Stored default_user_uuid is not a valid UUID — regenerating" + ); + let u = openfang_types::agent::UserId::new(); + memory + .structured_set( + shared, + KEY_UUID, + serde_json::Value::String(u.0.to_string()), + ) + .map_err(|e| { + KernelError::BootFailed(format!("kv write default_user_uuid: {e}")) + })?; + u + } + }, + _ => { + let u = openfang_types::agent::UserId::new(); + info!(default_user = %u.0, "Generated persistent default user UUID"); + memory + .structured_set(shared, KEY_UUID, serde_json::Value::String(u.0.to_string())) + .map_err(|e| { + KernelError::BootFailed(format!("kv write default_user_uuid: {e}")) + })?; + u + } + }; + + // Install for all callers of `default_user_id()`. + openfang_memory::session::set_default_user_id(user_id); + + // One-shot fixup: rewrite legacy nil-UUID sessions to the new default. + // + // The sentinel is written ONLY on success — if the rewrite fails + // (lock contention, disk error, etc.) the next boot will retry, + // rather than permanently leaving sessions stranded on the nil UUID. + let already_done = matches!( + memory.structured_get(shared, KEY_DONE).ok().flatten(), + Some(serde_json::Value::Bool(true)) + ); + if !already_done { + match memory.rewrite_nil_user_sessions(user_id) { + Ok(n) => { + if n > 0 { + info!( + sessions_updated = n, + "Rewrote nil-UUID sessions to persistent default user" + ); + } + // Mark the bootstrap done only after a clean rewrite — a + // failure path must NOT poison the sentinel, otherwise + // the retry would be silently suppressed on the next boot. + if let Err(e) = + memory.structured_set(shared, KEY_DONE, serde_json::Value::Bool(true)) + { + warn!(error = %e, "Failed to persist default_user_bootstrap_done sentinel"); + } + } + Err(e) => warn!( + error = %e, + "Failed to rewrite nil-UUID sessions; will retry on next boot" + ), + } + } + + Ok(()) + } } /// Convert a manifest's capability declarations into Capability enums. @@ -9412,4 +9632,231 @@ system_prompt = "You are a test agent." .expect("read pre-existing"); assert_eq!(contents, "hello", "must not overwrite user files"); } + + // ── PR 1: persistent default user + session user-tagging foundation ──────── + // + // These tests target the kernel-side bootstrap and the public/strict + // `resolve_user_id` filter. Each test boots a kernel against a temp dir, + // exercises the new APIs, and shuts down cleanly. + + fn minimal_kernel(tmp: &tempfile::TempDir) -> OpenFangKernel { + let home_dir = tmp.path().join("openfang"); + std::fs::create_dir_all(&home_dir).unwrap(); + let config = KernelConfig { + home_dir: home_dir.clone(), + data_dir: home_dir.join("data"), + ..KernelConfig::default() + }; + OpenFangKernel::boot_with_config(config).expect("kernel boots") + } + + /// Helper: read the persisted default-user UUID directly out of + /// `kv_store[shared_memory_agent_id, "default_user_uuid"]`. + /// + /// This bypasses `default_user_id()`, which reads a process-wide + /// `OnceLock` and would happily return the same value across restarts + /// even if `bootstrap_default_user` regenerated a fresh UUID on every + /// boot. Reading the durable storage is the only way to prove the boot + /// path actually persisted. + fn read_persisted_default_user_uuid(kernel: &OpenFangKernel) -> String { + let value = kernel + .memory + .structured_get(shared_memory_agent_id(), "default_user_uuid") + .expect("kv read default_user_uuid") + .expect("default_user_uuid must be persisted at boot"); + match value { + serde_json::Value::String(s) => s, + other => panic!("default_user_uuid stored as non-string: {other:?}"), + } + } + + /// The persistent default-user UUID must survive a kernel restart against + /// the same data directory. This is the "stable identity for one-person + /// installs" contract: every reboot resolves to the same UUID. + /// + /// We verify persistence by reading the durable `kv_store` row directly + /// — NOT via `default_user_id()`, which is backed by a process-wide + /// `OnceLock` and would return the same value across reboots even if + /// the bootstrap regenerated a fresh UUID on disk every time. + #[test] + fn test_bootstrap_default_user_persists_across_restarts() { + let tmp = tempfile::tempdir().unwrap(); + + // First boot: bootstrap generates and persists a new UUID. + let k1 = minimal_kernel(&tmp); + let first_uuid_str = read_persisted_default_user_uuid(&k1); + let first_uuid = uuid::Uuid::parse_str(&first_uuid_str) + .expect("persisted default_user_uuid must be a valid UUID"); + assert!( + !first_uuid.is_nil(), + "bootstrap must generate a non-nil UUID, got {first_uuid_str}" + ); + k1.shutdown(); + + // Second boot against the same data dir: must reuse the persisted + // UUID. Compare the raw kv_store strings — that is the only source + // of truth that survives a process restart. + let k2 = minimal_kernel(&tmp); + let second_uuid_str = read_persisted_default_user_uuid(&k2); + assert_eq!( + first_uuid_str, second_uuid_str, + "default_user_uuid in kv_store must persist across kernel restarts" + ); + k2.shutdown(); + } + + /// `resolve_user_id_internal` is the raw mapper used by in-process + /// callers. It must round-trip the test-user alias and pass arbitrary + /// UUIDs through unchanged so the integration suite can address the + /// test user. + #[test] + fn test_resolve_user_id_internal_preserves_test_alias() { + let tmp = tempfile::tempdir().unwrap(); + let kernel = minimal_kernel(&tmp); + + assert_eq!( + kernel.resolve_user_id_internal(Some("test")), + Some(openfang_types::agent::UserId( + openfang_memory::session::TEST_USER_UUID + )) + ); + + let uid = uuid::Uuid::new_v4(); + assert_eq!( + kernel.resolve_user_id_internal(Some(&uid.to_string())), + Some(openfang_types::agent::UserId(uid)) + ); + + // None and unparseable inputs both yield None. + assert_eq!(kernel.resolve_user_id_internal(None), None); + assert_eq!(kernel.resolve_user_id_internal(Some("not-a-uuid")), None); + + // "default" resolves to the persistent default user. + assert_eq!( + kernel.resolve_user_id_internal(Some("default")), + Some(kernel.default_user_id()) + ); + + kernel.shutdown(); + } + + /// The public `resolve_user_id` is the boundary between untrusted API + /// input and write-path attribution. It must fold both `"test"` and the + /// nil UUID to the persistent default user, leave other UUIDs alone, and + /// return `None` for unparseable input. + #[test] + fn test_resolve_user_id_strict_filter_folds_reserved_buckets() { + let tmp = tempfile::tempdir().unwrap(); + let kernel = minimal_kernel(&tmp); + let default = kernel.default_user_id(); + + // "test" → default (not the test user) + assert_eq!(kernel.resolve_user_id(Some("test")), Some(default)); + assert_ne!( + kernel.resolve_user_id(Some("test")), + Some(openfang_types::agent::UserId( + openfang_memory::session::TEST_USER_UUID + )) + ); + + // Nil UUID → default + let nil = uuid::Uuid::nil(); + assert_eq!( + kernel.resolve_user_id(Some(&nil.to_string())), + Some(default) + ); + + // Other valid UUIDs pass through. + let uid = uuid::Uuid::new_v4(); + assert_eq!( + kernel.resolve_user_id(Some(&uid.to_string())), + Some(openfang_types::agent::UserId(uid)) + ); + + // None / unparseable behave like the internal variant. + assert_eq!(kernel.resolve_user_id(None), None); + assert_eq!(kernel.resolve_user_id(Some("not-a-uuid")), None); + + // "default" still resolves to the default user. + assert_eq!(kernel.resolve_user_id(Some("default")), Some(default)); + + kernel.shutdown(); + } + + /// Sessions created without an explicit user_id are bound to the + /// persistent default user — the "single-user installs, everyone is the + /// default user" contract. + /// + /// We verify by reading the `user_id` column directly out of SQLite and + /// matching it against the argument passed to `create_session`, NOT + /// against `kernel.default_user_id()` returned at the assertion site. + /// Both reads-via-getter would tautologically agree on the OnceLock's + /// value even if the column was somehow stamped with a different one. + /// + /// We also verify that the persisted UUID is not nil and is a valid + /// UUID — those properties are insensitive to whichever kernel happened + /// to install the OnceLock first under cargo's shared-process test + /// harness. + #[test] + fn test_session_created_without_user_id_uses_default_user() { + let tmp = tempfile::tempdir().unwrap(); + let kernel = minimal_kernel(&tmp); + let agent_id = openfang_types::agent::AgentId::new(); + + // Capture the exact argument we'll pass to create_session so we can + // compare the column read-back to it (not to a later getter call). + let passed_user_id = openfang_memory::session::default_user_id(); + + let session = kernel + .memory + .create_session(agent_id, passed_user_id) + .expect("create session"); + + // Read the user_id column straight from SQLite — bypassing + // `SessionStore::get_session` which falls back to `default_user_id()` + // when the column is NULL, masking the very behaviour we want to + // confirm. + let conn_arc = kernel.memory.usage_conn(); + let conn = conn_arc.lock().unwrap(); + let stored_user_id: String = conn + .query_row( + "SELECT user_id FROM sessions WHERE id = ?1", + rusqlite::params![session.id.0.to_string()], + |row| row.get(0), + ) + .expect("session row must exist"); + drop(conn); + + let stored_uuid = + uuid::Uuid::parse_str(&stored_user_id).expect("user_id column must be a valid UUID"); + assert!( + !stored_uuid.is_nil(), + "default-user session must not be tagged with the nil UUID" + ); + assert_eq!( + stored_user_id, + passed_user_id.0.to_string(), + "user_id column must round-trip the value passed to create_session" + ); + + kernel.shutdown(); + } + + /// Sessions created with an explicit user_id are bound to that user. + #[test] + fn test_session_created_with_explicit_user_id_is_bound() { + let tmp = tempfile::tempdir().unwrap(); + let kernel = minimal_kernel(&tmp); + let agent_id = openfang_types::agent::AgentId::new(); + let alice = openfang_types::agent::UserId::new(); + + let session = kernel + .memory + .create_session(agent_id, alice) + .expect("create session"); + assert_eq!(session.user_id, alice); + assert_ne!(session.user_id, kernel.default_user_id()); + + kernel.shutdown(); + } } diff --git a/crates/openfang-memory/src/migration.rs b/crates/openfang-memory/src/migration.rs index 9249686510..88c64f77a3 100644 --- a/crates/openfang-memory/src/migration.rs +++ b/crates/openfang-memory/src/migration.rs @@ -5,7 +5,7 @@ use rusqlite::Connection; /// Current schema version. -const SCHEMA_VERSION: u32 = 8; +const SCHEMA_VERSION: u32 = 9; /// Run all migrations to bring the database up to date. pub fn run_migrations(conn: &Connection) -> Result<(), rusqlite::Error> { @@ -43,6 +43,10 @@ pub fn run_migrations(conn: &Connection) -> Result<(), rusqlite::Error> { migrate_v8(conn)?; } + if current_version < 9 { + migrate_v9(conn)?; + } + set_schema_version(conn, SCHEMA_VERSION)?; Ok(()) } @@ -328,6 +332,40 @@ fn migrate_v8(conn: &Connection) -> Result<(), rusqlite::Error> { Ok(()) } +/// Version 9: Tag every session with a user and a parent-session link. +/// +/// - `user_id` (TEXT, NOT NULL, default = nil UUID) — owning user. Existing +/// v8 rows inherit the nil-UUID sentinel; the kernel's +/// `bootstrap_default_user` rewrite migrates them to the persistent +/// default user on first boot after upgrade. +/// - `parent_session_id` (TEXT, nullable) — set when a session is forked off +/// another (e.g. a hand session linked back to its caller). No production +/// code path forks sessions yet, so the column is unused on existing rows. +/// +/// Indexes on `(agent_id, user_id)` and `parent_session_id` keep the per-user +/// lookups and child-session fan-out cheap. +fn migrate_v9(conn: &Connection) -> Result<(), rusqlite::Error> { + if !column_exists(conn, "sessions", "user_id") { + conn.execute( + "ALTER TABLE sessions ADD COLUMN user_id TEXT NOT NULL DEFAULT '00000000-0000-0000-0000-000000000000'", + [], + )?; + } + if !column_exists(conn, "sessions", "parent_session_id") { + conn.execute("ALTER TABLE sessions ADD COLUMN parent_session_id TEXT", [])?; + } + conn.execute_batch( + " + CREATE INDEX IF NOT EXISTS idx_sessions_user ON sessions(agent_id, user_id); + CREATE INDEX IF NOT EXISTS idx_sessions_parent ON sessions(parent_session_id); + + INSERT OR IGNORE INTO migrations (version, applied_at, description) + VALUES (9, datetime('now'), 'Tag sessions with user_id and parent_session_id'); + ", + )?; + Ok(()) +} + #[cfg(test)] mod tests { use super::*; @@ -360,4 +398,78 @@ mod tests { run_migrations(&conn).unwrap(); run_migrations(&conn).unwrap(); // Should not error } + + #[test] + fn test_migration_v9_adds_session_user_columns() { + let conn = Connection::open_in_memory().unwrap(); + run_migrations(&conn).unwrap(); + assert!(column_exists(&conn, "sessions", "user_id")); + assert!(column_exists(&conn, "sessions", "parent_session_id")); + } + + #[test] + fn test_migration_v9_indexes_present() { + // The (agent_id, user_id) and parent_session_id indexes underpin the + // per-user and child-session lookups added in PR 1. If either is + // missing, the queries fall back to full table scans. + let conn = Connection::open_in_memory().unwrap(); + run_migrations(&conn).unwrap(); + let names: Vec = conn + .prepare("SELECT name FROM sqlite_master WHERE type='index'") + .unwrap() + .query_map([], |row| row.get(0)) + .unwrap() + .filter_map(|r| r.ok()) + .collect(); + assert!(names.contains(&"idx_sessions_user".to_string())); + assert!(names.contains(&"idx_sessions_parent".to_string())); + } + + /// A v8 database (sessions table without `user_id` or `parent_session_id`) + /// must upgrade cleanly to v9: existing rows survive, the new columns are + /// present, and the user_id default sentinel is the nil UUID. + #[test] + fn test_migration_v8_to_v9_upgrade_preserves_existing_rows() { + let conn = Connection::open_in_memory().unwrap(); + + // Build the v8 schema by stopping the migration runner one step early. + migrate_v1(&conn).unwrap(); + migrate_v2(&conn).unwrap(); + migrate_v3(&conn).unwrap(); + migrate_v4(&conn).unwrap(); + migrate_v5(&conn).unwrap(); + migrate_v6(&conn).unwrap(); + migrate_v7(&conn).unwrap(); + migrate_v8(&conn).unwrap(); + set_schema_version(&conn, 8).unwrap(); + + assert!(!column_exists(&conn, "sessions", "user_id")); + assert!(!column_exists(&conn, "sessions", "parent_session_id")); + + // Insert a pre-v9 session row (no user_id, no parent_session_id). + conn.execute( + "INSERT INTO sessions (id, agent_id, messages, context_window_tokens, created_at, updated_at) + VALUES ('sess-1', 'agent-1', X'', 0, datetime('now'), datetime('now'))", + [], + ) + .unwrap(); + + // Now run the full pipeline — v9 should apply on top of v8. + run_migrations(&conn).unwrap(); + + assert!(column_exists(&conn, "sessions", "user_id")); + assert!(column_exists(&conn, "sessions", "parent_session_id")); + + // The pre-existing row survives and gets the nil-UUID default for + // user_id, NULL for parent_session_id. + let (uid, parent): (String, Option) = conn + .query_row( + "SELECT user_id, parent_session_id FROM sessions WHERE id = 'sess-1'", + [], + |row| Ok((row.get::<_, String>(0)?, row.get::<_, Option>(1)?)), + ) + .unwrap(); + assert_eq!(uid, "00000000-0000-0000-0000-000000000000"); + assert!(parent.is_none()); + } } diff --git a/crates/openfang-memory/src/session.rs b/crates/openfang-memory/src/session.rs index 7a18e0f9fb..1176a58e50 100644 --- a/crates/openfang-memory/src/session.rs +++ b/crates/openfang-memory/src/session.rs @@ -1,13 +1,70 @@ //! Session management — load/save conversation history. use chrono::Utc; -use openfang_types::agent::{AgentId, SessionId}; +use openfang_types::agent::{AgentId, SessionId, UserId}; use openfang_types::error::{OpenFangError, OpenFangResult}; use openfang_types::message::{ContentBlock, Message, MessageContent, Role}; use rusqlite::Connection; use std::io::Write; use std::path::Path; -use std::sync::{Arc, Mutex}; +use std::sync::{Arc, Mutex, OnceLock}; + +/// Process-wide override for the default user ID. +/// +/// Set by the kernel at boot via [`set_default_user_id`] after it loads or +/// generates the persistent default-user UUID from the kv_store. If unset, +/// [`default_user_id`] falls back to the legacy nil UUID so library callers +/// (tests, embeddings without a kernel) still get a deterministic value. +static DEFAULT_USER_ID: OnceLock = OnceLock::new(); + +/// Install the persistent default user ID. Called once at kernel boot. +/// +/// Subsequent calls are ignored (OnceLock semantics). This matches the +/// expectation that a single kernel process has exactly one default user. +pub fn set_default_user_id(user_id: UserId) { + let _ = DEFAULT_USER_ID.set(user_id); +} + +/// Return the default user ID. +/// +/// When the kernel has installed a persistent UUID via [`set_default_user_id`] +/// (the normal case in production), returns it. Otherwise falls back to the +/// nil UUID — preserved for tests and embedded use that never call the kernel. +pub fn default_user_id() -> UserId { + DEFAULT_USER_ID + .get() + .copied() + .unwrap_or_else(|| UserId(uuid::Uuid::nil())) +} + +/// Sentinel UUID for the well-known "test user" bucket +/// (`00000000-0000-0000-0000-000000000002`). +/// +/// This is distinct from: +/// - the persistent default user (a fresh UUID generated at kernel boot, +/// falling back to the nil UUID `…0000` for library callers that never +/// call the kernel), and +/// - `shared_memory_agent_id()`, which uses +/// `00000000-0000-0000-0000-000000000001` as its agent-side sentinel. +/// +/// Production code may reference this constant only to recognise the +/// kernel's `"test"` alias and (for the strict HTTP-boundary resolver) fold +/// it back to the persistent default user. It is **not** routable from any +/// external API surface — see `OpenFangKernel::resolve_user_id`. +pub const TEST_USER_UUID: uuid::Uuid = uuid::Uuid::from_u128(2); + +/// Return the test user ID — a fixed well-known UUID separate from the +/// default user and from `shared_memory_agent_id()`. +/// +/// Test-only helper used by the in-process integration suite to address an +/// isolated bucket without polluting the default user's memory. Production +/// code that needs the same UUID for sentinel comparisons must reference +/// [`TEST_USER_UUID`] directly so this helper stays out of the production +/// API surface. +#[cfg(test)] +pub fn test_user_id() -> UserId { + UserId(TEST_USER_UUID) +} /// A conversation session with message history. #[derive(Debug, Clone)] @@ -16,6 +73,17 @@ pub struct Session { pub id: SessionId, /// Owning agent ID. pub agent_id: AgentId, + /// User this session belongs to. + /// + /// Required: every session is owned by exactly one user. New sessions + /// created without an explicit user use [`default_user_id`]. + pub user_id: UserId, + /// Parent session ID — set when this session was forked off another one + /// (e.g. a hand session linked back to its caller). Lays the foundation + /// for tree-scoped cascade deletion; no production code path forks + /// sessions in this PR, so the value is `None` for every session created + /// today. + pub parent_session_id: Option, /// Conversation messages. pub messages: Vec, /// Estimated token count for the context window. @@ -43,7 +111,11 @@ impl SessionStore { .lock() .map_err(|e| OpenFangError::Internal(e.to_string()))?; let mut stmt = conn - .prepare("SELECT agent_id, messages, context_window_tokens, label FROM sessions WHERE id = ?1") + .prepare( + "SELECT agent_id, messages, context_window_tokens, label, \ + user_id, parent_session_id \ + FROM sessions WHERE id = ?1", + ) .map_err(|e| OpenFangError::Memory(e.to_string()))?; let result = stmt.query_row(rusqlite::params![session_id.0.to_string()], |row| { @@ -51,19 +123,37 @@ impl SessionStore { let messages_blob: Vec = row.get(1)?; let tokens: i64 = row.get(2)?; let label: Option = row.get(3).unwrap_or(None); - Ok((agent_str, messages_blob, tokens, label)) + let user_id_str: Option = row.get(4).unwrap_or(None); + let parent_session_id_str: Option = row.get(5).unwrap_or(None); + Ok(( + agent_str, + messages_blob, + tokens, + label, + user_id_str, + parent_session_id_str, + )) }); match result { - Ok((agent_str, messages_blob, tokens, label)) => { + Ok((agent_str, messages_blob, tokens, label, user_id_str, parent_session_id_str)) => { let agent_id = uuid::Uuid::parse_str(&agent_str) .map(AgentId) .map_err(|e| OpenFangError::Memory(e.to_string()))?; let messages: Vec = rmp_serde::from_slice(&messages_blob) .map_err(|e| OpenFangError::Serialization(e.to_string()))?; + let user_id = user_id_str + .and_then(|s| uuid::Uuid::parse_str(&s).ok()) + .map(UserId) + .unwrap_or_else(default_user_id); + let parent_session_id = parent_session_id_str + .and_then(|s| uuid::Uuid::parse_str(&s).ok()) + .map(SessionId); Ok(Some(Session { id: session_id, agent_id, + user_id, + parent_session_id, messages, context_window_tokens: tokens as u64, label, @@ -83,16 +173,26 @@ impl SessionStore { let messages_blob = rmp_serde::to_vec_named(&session.messages) .map_err(|e| OpenFangError::Serialization(e.to_string()))?; let now = Utc::now().to_rfc3339(); + let parent_str = session + .parent_session_id + .as_ref() + .map(|id| id.0.to_string()); conn.execute( - "INSERT INTO sessions (id, agent_id, messages, context_window_tokens, label, created_at, updated_at) - VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?6) - ON CONFLICT(id) DO UPDATE SET messages = ?3, context_window_tokens = ?4, label = ?5, updated_at = ?6", + "INSERT INTO sessions \ + (id, agent_id, messages, context_window_tokens, label, \ + user_id, parent_session_id, created_at, updated_at) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?8) + ON CONFLICT(id) DO UPDATE SET \ + messages = ?3, context_window_tokens = ?4, label = ?5, \ + user_id = ?6, parent_session_id = ?7, updated_at = ?8", rusqlite::params![ session.id.0.to_string(), session.agent_id.0.to_string(), messages_blob, session.context_window_tokens as i64, session.label.as_deref(), + session.user_id.0.to_string(), + parent_str, now, ], ) @@ -182,11 +282,16 @@ impl SessionStore { Ok(sessions) } - /// Create a new empty session for an agent. - pub fn create_session(&self, agent_id: AgentId) -> OpenFangResult { + /// Create a new empty session for an agent owned by `user_id`. + /// + /// Callers that have no specific user pass [`default_user_id`] so the + /// session attaches to the kernel's persistent default identity. + pub fn create_session(&self, agent_id: AgentId, user_id: UserId) -> OpenFangResult { let session = Session { id: SessionId::new(), agent_id, + user_id, + parent_session_id: None, messages: Vec::new(), context_window_tokens: 0, label: None, @@ -225,8 +330,9 @@ impl SessionStore { .map_err(|e| OpenFangError::Internal(e.to_string()))?; let mut stmt = conn .prepare( - "SELECT id, messages, context_window_tokens, label FROM sessions \ - WHERE agent_id = ?1 AND label = ?2 LIMIT 1", + "SELECT id, messages, context_window_tokens, label, \ + user_id, parent_session_id \ + FROM sessions WHERE agent_id = ?1 AND label = ?2 LIMIT 1", ) .map_err(|e| OpenFangError::Memory(e.to_string()))?; @@ -235,19 +341,37 @@ impl SessionStore { let messages_blob: Vec = row.get(1)?; let tokens: i64 = row.get(2)?; let lbl: Option = row.get(3).unwrap_or(None); - Ok((id_str, messages_blob, tokens, lbl)) + let user_id_str: Option = row.get(4).unwrap_or(None); + let parent_session_id_str: Option = row.get(5).unwrap_or(None); + Ok(( + id_str, + messages_blob, + tokens, + lbl, + user_id_str, + parent_session_id_str, + )) }); match result { - Ok((id_str, messages_blob, tokens, lbl)) => { + Ok((id_str, messages_blob, tokens, lbl, user_id_str, parent_session_id_str)) => { let session_id = uuid::Uuid::parse_str(&id_str) .map(SessionId) .map_err(|e| OpenFangError::Memory(e.to_string()))?; let messages: Vec = rmp_serde::from_slice(&messages_blob) .map_err(|e| OpenFangError::Serialization(e.to_string()))?; + let user_id = user_id_str + .and_then(|s| uuid::Uuid::parse_str(&s).ok()) + .map(UserId) + .unwrap_or_else(default_user_id); + let parent_session_id = parent_session_id_str + .and_then(|s| uuid::Uuid::parse_str(&s).ok()) + .map(SessionId); Ok(Some(Session { id: session_id, agent_id, + user_id, + parent_session_id, messages, context_window_tokens: tokens as u64, label: lbl, @@ -298,6 +422,10 @@ impl SessionStore { } /// Create a new session with an optional label. + /// + /// The session is owned by [`default_user_id`]. This entry point is used + /// by the kernel's "create named session" API, which today does not + /// surface a user selector. pub fn create_session_with_label( &self, agent_id: AgentId, @@ -306,6 +434,8 @@ impl SessionStore { let session = Session { id: SessionId::new(), agent_id, + user_id: default_user_id(), + parent_session_id: None, messages: Vec::new(), context_window_tokens: 0, label: label.map(|s| s.to_string()), @@ -635,18 +765,73 @@ mod tests { fn test_create_and_load_session() { let store = setup(); let agent_id = AgentId::new(); - let session = store.create_session(agent_id).unwrap(); + let session = store.create_session(agent_id, default_user_id()).unwrap(); let loaded = store.get_session(session.id).unwrap().unwrap(); assert_eq!(loaded.agent_id, agent_id); assert!(loaded.messages.is_empty()); + // Sessions created with the default user must report that ownership + // through the new column. + assert_eq!(loaded.user_id, default_user_id()); + assert!(loaded.parent_session_id.is_none()); + } + + #[test] + fn test_create_session_with_explicit_user_binds_to_that_user() { + // Sessions created with a specific user_id must round-trip through + // SQLite without being silently rewritten to the default user. + let store = setup(); + let agent_id = AgentId::new(); + let user = UserId::new(); + + let session = store.create_session(agent_id, user).unwrap(); + let loaded = store.get_session(session.id).unwrap().unwrap(); + assert_eq!(loaded.user_id, user); + assert_ne!(loaded.user_id, default_user_id()); + } + + #[test] + fn test_default_and_test_user_ids_differ() { + // `default_user_id` and `test_user_id` must be distinct so the test + // bucket never collides with the production default identity. + // (OnceLock is process-wide so we don't assert the nil-UUID fallback + // value directly — another test in the binary may have installed a + // persistent UUID first.) + assert_ne!(default_user_id(), test_user_id()); + assert_eq!(test_user_id().0, TEST_USER_UUID); + assert_eq!(TEST_USER_UUID, uuid::Uuid::from_u128(2)); + } + + #[test] + fn test_user_id_does_not_collide_with_shared_memory_agent() { + // `shared_memory_agent_id()` lives on the agent-ID axis and uses + // `from_u128(1)`. The test user bucket uses `from_u128(2)` so they + // can never accidentally compare equal when a stringly-typed UUID + // appears in logs or migration audits. + let shared_agent_uuid = uuid::Uuid::from_u128(1); + assert_ne!(TEST_USER_UUID, shared_agent_uuid); + } + + #[test] + fn test_save_session_persists_parent_link() { + // Sessions with a `parent_session_id` round-trip the link so a + // future PR can cascade-delete by session tree. + let store = setup(); + let agent_id = AgentId::new(); + let parent = store.create_session(agent_id, default_user_id()).unwrap(); + let mut child = store.create_session(agent_id, default_user_id()).unwrap(); + child.parent_session_id = Some(parent.id); + store.save_session(&child).unwrap(); + + let loaded = store.get_session(child.id).unwrap().unwrap(); + assert_eq!(loaded.parent_session_id, Some(parent.id)); } #[test] fn test_save_and_load_with_messages() { let store = setup(); let agent_id = AgentId::new(); - let mut session = store.create_session(agent_id).unwrap(); + let mut session = store.create_session(agent_id, default_user_id()).unwrap(); session.messages.push(Message::user("Hello")); session.messages.push(Message::assistant("Hi there!")); store.save_session(&session).unwrap(); @@ -666,7 +851,7 @@ mod tests { fn test_delete_session() { let store = setup(); let agent_id = AgentId::new(); - let session = store.create_session(agent_id).unwrap(); + let session = store.create_session(agent_id, default_user_id()).unwrap(); let sid = session.id; assert!(store.get_session(sid).unwrap().is_some()); store.delete_session(sid).unwrap(); @@ -677,8 +862,8 @@ mod tests { fn test_delete_agent_sessions() { let store = setup(); let agent_id = AgentId::new(); - let s1 = store.create_session(agent_id).unwrap(); - let s2 = store.create_session(agent_id).unwrap(); + let s1 = store.create_session(agent_id, default_user_id()).unwrap(); + let s2 = store.create_session(agent_id, default_user_id()).unwrap(); assert!(store.get_session(s1.id).unwrap().is_some()); assert!(store.get_session(s2.id).unwrap().is_some()); store.delete_agent_sessions(agent_id).unwrap(); @@ -782,7 +967,7 @@ mod tests { fn test_jsonl_mirror_write() { let store = setup(); let agent_id = AgentId::new(); - let mut session = store.create_session(agent_id).unwrap(); + let mut session = store.create_session(agent_id, default_user_id()).unwrap(); session .messages .push(openfang_types::message::Message::user("Hello")); diff --git a/crates/openfang-memory/src/substrate.rs b/crates/openfang-memory/src/substrate.rs index 03ecd78cfc..41dc160bba 100644 --- a/crates/openfang-memory/src/substrate.rs +++ b/crates/openfang-memory/src/substrate.rs @@ -12,7 +12,7 @@ use crate::structured::StructuredStore; use crate::usage::UsageStore; use async_trait::async_trait; -use openfang_types::agent::{AgentEntry, AgentId, SessionId}; +use openfang_types::agent::{AgentEntry, AgentId, SessionId, UserId}; use openfang_types::config::MemoryConfig; use openfang_types::error::{OpenFangError, OpenFangResult}; use openfang_types::memory::{ @@ -207,9 +207,51 @@ impl MemorySubstrate { .map_err(|e| OpenFangError::Internal(e.to_string()))? } - /// Create a new empty session for an agent. - pub fn create_session(&self, agent_id: AgentId) -> OpenFangResult { - self.sessions.create_session(agent_id) + /// Create a new empty session for an agent owned by `user_id`. + /// + /// Callers without a specific user pass `default_user_id()` so the + /// session attaches to the kernel's persistent default identity. + pub fn create_session(&self, agent_id: AgentId, user_id: UserId) -> OpenFangResult { + self.sessions.create_session(agent_id, user_id) + } + + /// Rewrite all sessions that are still tagged with the nil-UUID owner + /// (the legacy "anonymous bucket" default) to the given `target` user. + /// + /// Returns the number of sessions updated. Called once at kernel boot + /// after the persistent default-user UUID is determined; subsequent calls + /// are no-ops because no nil-UUID rows remain. + /// + /// The UPDATE runs inside an explicit transaction so a crash or lock + /// failure mid-flight rolls back cleanly — the caller's bootstrap + /// sentinel relies on this being all-or-nothing to avoid permanently + /// stranding sessions on the nil bucket. + pub fn rewrite_nil_user_sessions(&self, target: UserId) -> OpenFangResult { + let mut conn = self + .conn + .lock() + .map_err(|e| OpenFangError::Internal(e.to_string()))?; + let nil_str = uuid::Uuid::nil().to_string(); + let target_str = target.0.to_string(); + + let tx = conn + .transaction() + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + let sessions_updated = tx + .execute( + "UPDATE sessions SET user_id = ?1 WHERE user_id = ?2 OR user_id IS NULL", + rusqlite::params![target_str, nil_str], + ) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + tx.commit() + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + + info!( + sessions_updated = sessions_updated, + target = %target.0, + "Rewrote nil-UUID sessions to persistent default user" + ); + Ok(sessions_updated) } /// List all sessions with metadata. @@ -731,6 +773,109 @@ impl Memory for MemorySubstrate { #[cfg(test)] mod tests { use super::*; + use crate::session::default_user_id; + + #[test] + fn test_rewrite_nil_user_sessions_is_idempotent_and_targeted() { + // The one-shot fixup must rewrite ONLY the nil-UUID rows and leave + // every other session alone; a second call after the rewrite must + // find zero matching rows and behave as a no-op. + let substrate = MemorySubstrate::open_in_memory(0.1).unwrap(); + let agent_id = AgentId::new(); + let other_user = UserId::new(); + let target = UserId::new(); + let nil_user = UserId(uuid::Uuid::nil()); + + // Two nil-bucket sessions (the legacy default) + one explicit-user session. + let s1 = substrate + .sessions + .create_session(agent_id, nil_user) + .unwrap(); + let s2 = substrate + .sessions + .create_session(agent_id, nil_user) + .unwrap(); + let s3 = substrate + .sessions + .create_session(agent_id, other_user) + .unwrap(); + + let updated = substrate.rewrite_nil_user_sessions(target).unwrap(); + assert_eq!( + updated, 2, + "exactly the two nil-UUID sessions must be rewritten" + ); + + let s1 = substrate.sessions.get_session(s1.id).unwrap().unwrap(); + let s2 = substrate.sessions.get_session(s2.id).unwrap().unwrap(); + let s3 = substrate.sessions.get_session(s3.id).unwrap().unwrap(); + assert_eq!(s1.user_id, target); + assert_eq!(s2.user_id, target); + assert_eq!( + s3.user_id, other_user, + "non-nil sessions must NOT be touched" + ); + + // Second call is a no-op (no nil rows left). + let updated2 = substrate.rewrite_nil_user_sessions(target).unwrap(); + assert_eq!(updated2, 0); + } + + #[test] + fn test_rewrite_nil_user_sessions_returns_err_when_table_missing() { + // When the `sessions` table is missing, the UPDATE cannot run and the + // caller must observe an error — not a silent success. This is what + // lets `bootstrap_default_user` skip setting the + // `default_user_bootstrap_done` sentinel on failure so the rewrite is + // retried on the next boot. + // + // NOTE: This test does NOT exercise transactional atomicity. A single + // `UPDATE` is implicitly all-or-nothing under SQLite, so there is no + // partial-commit state to roll back from. True multi-statement + // atomicity becomes testable in PR 2 when the `extractions` table + // joins the same transaction — at that point we add a dedicated + // `…_is_atomic_on_partial_failure` test that forces the second + // statement to fail and verifies the first one is rolled back. + let substrate = MemorySubstrate::open_in_memory(0.1).unwrap(); + let agent_id = AgentId::new(); + let target = UserId::new(); + let nil_user = UserId(uuid::Uuid::nil()); + + let _s1 = substrate + .sessions + .create_session(agent_id, nil_user) + .unwrap(); + + // Drop the only table the UPDATE touches; the next call must error. + { + let conn = substrate.conn.lock().unwrap(); + conn.execute("DROP TABLE sessions", []) + .expect("drop sessions"); + } + + let err = substrate.rewrite_nil_user_sessions(target); + assert!( + err.is_err(), + "rewrite must fail when the UPDATE cannot complete" + ); + } + + #[test] + fn test_create_session_via_substrate_routes_user_id() { + // The substrate's `create_session` proxy must forward the user_id + // through to the session store rather than silently dropping it. + let substrate = MemorySubstrate::open_in_memory(0.1).unwrap(); + let agent_id = AgentId::new(); + let user = UserId::new(); + + let session = substrate.create_session(agent_id, user).unwrap(); + assert_eq!(session.user_id, user); + + let default_session = substrate + .create_session(agent_id, default_user_id()) + .unwrap(); + assert_eq!(default_session.user_id, default_user_id()); + } #[tokio::test] async fn test_substrate_kv() { diff --git a/crates/openfang-runtime/src/agent_loop.rs b/crates/openfang-runtime/src/agent_loop.rs index 615fbfdb73..6f29244fc6 100644 --- a/crates/openfang-runtime/src/agent_loop.rs +++ b/crates/openfang-runtime/src/agent_loop.rs @@ -3792,6 +3792,8 @@ mod tests { let mut session = openfang_memory::session::Session { id: openfang_types::agent::SessionId::new(), agent_id, + user_id: openfang_memory::session::default_user_id(), + parent_session_id: None, messages: Vec::new(), context_window_tokens: 0, label: None, @@ -3845,6 +3847,8 @@ mod tests { let mut session = openfang_memory::session::Session { id: openfang_types::agent::SessionId::new(), agent_id, + user_id: openfang_memory::session::default_user_id(), + parent_session_id: None, messages: Vec::new(), context_window_tokens: 0, label: None, @@ -3900,6 +3904,8 @@ mod tests { let mut session = openfang_memory::session::Session { id: openfang_types::agent::SessionId::new(), agent_id, + user_id: openfang_memory::session::default_user_id(), + parent_session_id: None, messages: Vec::new(), context_window_tokens: 0, label: None, @@ -3953,6 +3959,8 @@ mod tests { let mut session = openfang_memory::session::Session { id: openfang_types::agent::SessionId::new(), agent_id, + user_id: openfang_memory::session::default_user_id(), + parent_session_id: None, messages: Vec::new(), context_window_tokens: 0, label: None, @@ -3997,6 +4005,8 @@ mod tests { let mut session = openfang_memory::session::Session { id: openfang_types::agent::SessionId::new(), agent_id, + user_id: openfang_memory::session::default_user_id(), + parent_session_id: None, messages: Vec::new(), context_window_tokens: 0, label: None, @@ -4123,6 +4133,8 @@ mod tests { let mut session = openfang_memory::session::Session { id: openfang_types::agent::SessionId::new(), agent_id, + user_id: openfang_memory::session::default_user_id(), + parent_session_id: None, messages: Vec::new(), context_window_tokens: 0, label: None, @@ -4170,6 +4182,8 @@ mod tests { let mut session = openfang_memory::session::Session { id: openfang_types::agent::SessionId::new(), agent_id, + user_id: openfang_memory::session::default_user_id(), + parent_session_id: None, messages: Vec::new(), context_window_tokens: 0, label: None, @@ -4223,6 +4237,8 @@ mod tests { let mut session = openfang_memory::session::Session { id: openfang_types::agent::SessionId::new(), agent_id, + user_id: openfang_memory::session::default_user_id(), + parent_session_id: None, messages: Vec::new(), context_window_tokens: 0, label: None, @@ -5189,6 +5205,8 @@ mod tests { let mut session = openfang_memory::session::Session { id: openfang_types::agent::SessionId::new(), agent_id, + user_id: openfang_memory::session::default_user_id(), + parent_session_id: None, messages: Vec::new(), context_window_tokens: 0, label: None, @@ -5260,6 +5278,8 @@ mod tests { let mut session = openfang_memory::session::Session { id: openfang_types::agent::SessionId::new(), agent_id, + user_id: openfang_memory::session::default_user_id(), + parent_session_id: None, messages: Vec::new(), context_window_tokens: 0, label: None, @@ -5337,6 +5357,8 @@ mod tests { let mut session = openfang_memory::session::Session { id: openfang_types::agent::SessionId::new(), agent_id, + user_id: openfang_memory::session::default_user_id(), + parent_session_id: None, messages: Vec::new(), context_window_tokens: 0, label: None, @@ -5392,6 +5414,8 @@ mod tests { let mut session = openfang_memory::session::Session { id: openfang_types::agent::SessionId::new(), agent_id, + user_id: openfang_memory::session::default_user_id(), + parent_session_id: None, messages: Vec::new(), context_window_tokens: 0, label: None, diff --git a/crates/openfang-runtime/src/compactor.rs b/crates/openfang-runtime/src/compactor.rs index 3602ff244a..2b50a236c1 100644 --- a/crates/openfang-runtime/src/compactor.rs +++ b/crates/openfang-runtime/src/compactor.rs @@ -777,6 +777,8 @@ mod tests { let session = Session { id: openfang_types::agent::SessionId::new(), agent_id: openfang_types::agent::AgentId::new(), + user_id: openfang_memory::session::default_user_id(), + parent_session_id: None, messages: vec![Message::user("hello")], context_window_tokens: 0, label: None, @@ -793,6 +795,8 @@ mod tests { let session = Session { id: openfang_types::agent::SessionId::new(), agent_id: openfang_types::agent::AgentId::new(), + user_id: openfang_memory::session::default_user_id(), + parent_session_id: None, messages, context_window_tokens: 0, label: None, @@ -842,6 +846,8 @@ mod tests { let session = Session { id: openfang_types::agent::SessionId::new(), agent_id: openfang_types::agent::AgentId::new(), + user_id: openfang_memory::session::default_user_id(), + parent_session_id: None, messages: vec![Message::user("hello"), Message::assistant("hi")], context_window_tokens: 0, label: None, @@ -931,6 +937,8 @@ mod tests { let session = Session { id: openfang_types::agent::SessionId::new(), agent_id: openfang_types::agent::AgentId::new(), + user_id: openfang_memory::session::default_user_id(), + parent_session_id: None, messages, context_window_tokens: 0, label: None, @@ -1002,6 +1010,8 @@ mod tests { let session = Session { id: openfang_types::agent::SessionId::new(), agent_id: openfang_types::agent::AgentId::new(), + user_id: openfang_memory::session::default_user_id(), + parent_session_id: None, messages, context_window_tokens: 0, label: None, @@ -1129,6 +1139,8 @@ mod tests { let session = Session { id: openfang_types::agent::SessionId::new(), agent_id: openfang_types::agent::AgentId::new(), + user_id: openfang_memory::session::default_user_id(), + parent_session_id: None, messages, context_window_tokens: 0, label: None, diff --git a/crates/openfang-types/src/config.rs b/crates/openfang-types/src/config.rs index 25df9b0059..a7e0fe23a7 100644 --- a/crates/openfang-types/src/config.rs +++ b/crates/openfang-types/src/config.rs @@ -186,6 +186,13 @@ pub struct UserConfig { /// Optional API key hash for API authentication. #[serde(default)] pub api_key_hash: Option, + /// Marks this user as the persistent default. When `true`, the kernel + /// binds this user's display name and channel bindings to the persistent + /// default-user UUID generated at boot rather than a freshly-generated + /// per-config UUID. At most one `[[users]]` block should set this; if + /// none does, the first user in the config inherits the default identity. + #[serde(default)] + pub is_default: bool, } fn default_role() -> String { @@ -4120,12 +4127,44 @@ mod tests { m }, api_key_hash: None, + is_default: false, }; let json = serde_json::to_string(&uc).unwrap(); let back: UserConfig = serde_json::from_str(&json).unwrap(); assert_eq!(back.name, "Alice"); assert_eq!(back.role, "owner"); assert_eq!(back.channel_bindings.get("telegram").unwrap(), "123456"); + assert!(!back.is_default); + } + + #[test] + fn test_user_config_is_default_defaults_to_false() { + // TOML/JSON without `is_default` must deserialize as `false` so that + // existing configs are unaffected by this PR. + let toml_src = r#" + name = "Alice" + role = "owner" + "#; + let uc: UserConfig = toml::from_str(toml_src).unwrap(); + assert!(!uc.is_default); + + let json_src = r#"{"name":"Bob","role":"user"}"#; + let uc: UserConfig = serde_json::from_str(json_src).unwrap(); + assert!(!uc.is_default); + } + + #[test] + fn test_user_config_is_default_roundtrips_when_set() { + let uc = UserConfig { + name: "Owner".to_string(), + role: "owner".to_string(), + channel_bindings: std::collections::HashMap::new(), + api_key_hash: None, + is_default: true, + }; + let json = serde_json::to_string(&uc).unwrap(); + let back: UserConfig = serde_json::from_str(&json).unwrap(); + assert!(back.is_default); } #[test] diff --git a/crates/openfang-types/src/message.rs b/crates/openfang-types/src/message.rs index ba5e93438d..e798b7108c 100644 --- a/crates/openfang-types/src/message.rs +++ b/crates/openfang-types/src/message.rs @@ -13,6 +13,20 @@ fn new_msg_id() -> String { uuid::Uuid::new_v4().to_string() } +/// Source tag for a message — identifies system-injected context that should +/// not be treated like ordinary conversational turns by downstream consumers +/// (e.g. structured-memory extraction, dream summarisation). +/// +/// Foundational type: this PR adds the field so older sessions deserialize +/// cleanly and future consumers (PR 2: per-user memory extraction filtering) +/// have a stable tag to read. No code path in this PR sets a non-`None` value. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum MessageSource { + /// Injected by context sources (e.g. calendar-hand, mail-hand summaries). + ContextInjection, +} + /// A message in an LLM conversation. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Message { @@ -38,6 +52,12 @@ pub struct Message { pub role: Role, /// The content of the message. pub content: MessageContent, + /// Optional source tag — set for system-injected messages that downstream + /// consumers may want to treat differently (e.g. exclude from structured + /// extraction). `#[serde(default)]` so sessions persisted before this + /// field existed deserialize with `source = None`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub source: Option, } /// The role of a message sender in an LLM conversation. @@ -76,6 +96,7 @@ impl Default for Message { provider_msg_id: None, role: Role::User, content: MessageContent::Text(String::new()), + source: None, } } } @@ -249,6 +270,7 @@ impl Message { provider_msg_id: None, role: Role::System, content: MessageContent::Text(content.into()), + source: None, } } @@ -259,6 +281,7 @@ impl Message { provider_msg_id: None, role: Role::User, content: MessageContent::Text(content.into()), + source: None, } } @@ -269,6 +292,7 @@ impl Message { provider_msg_id: None, role: Role::User, content: MessageContent::Blocks(blocks), + source: None, } } @@ -279,6 +303,7 @@ impl Message { provider_msg_id: None, role: Role::Assistant, content: MessageContent::Text(content.into()), + source: None, } } @@ -292,6 +317,7 @@ impl Message { provider_msg_id: None, role: Role::Assistant, content: MessageContent::Blocks(blocks), + source: None, } } @@ -587,6 +613,65 @@ mod tests { assert_eq!(restored.provider_msg_id.as_deref(), Some("msg_xyz")); } + #[test] + fn test_message_source_roundtrip() { + // The `MessageSource::ContextInjection` tag must survive serde + // round-trips so future consumers can rely on it. + let mut msg = Message::user("Hello"); + msg.source = Some(MessageSource::ContextInjection); + let json = serde_json::to_string(&msg).unwrap(); + let decoded: Message = serde_json::from_str(&json).unwrap(); + assert_eq!(decoded.source, Some(MessageSource::ContextInjection)); + } + + #[test] + fn test_message_source_skipped_when_none() { + // `skip_serializing_if = "Option::is_none"` keeps the JSON minimal + // for the overwhelmingly common case of an untagged user/assistant + // message. + let msg = Message::user("Hello"); + let json = serde_json::to_value(&msg).unwrap(); + assert!( + json.get("source").is_none(), + "untagged messages must not write a `source` field" + ); + } + + #[test] + fn test_message_source_legacy_deser_defaults_to_none() { + // A message serialized before the `source` field existed should + // deserialize with `source = None`. This is what every existing + // session on disk looks like — `serde(default)` carries the load. + let json = r#"{"role":"user","content":"Hello"}"#; + let msg: Message = serde_json::from_str(json).unwrap(); + assert_eq!(msg.source, None); + assert_eq!(msg.role, Role::User); + } + + #[test] + fn test_message_source_msgpack_roundtrip() { + // SessionStore persists messages via rmp_serde::to_vec_named — make + // sure the new field round-trips through msgpack the same way it + // does through JSON, since msgpack is the on-disk format. + let mut msg = Message::user("hello"); + msg.source = Some(MessageSource::ContextInjection); + let bytes = rmp_serde::to_vec_named(&msg).expect("msgpack encode"); + let restored: Message = rmp_serde::from_slice(&bytes).expect("msgpack decode"); + assert_eq!(restored.source, Some(MessageSource::ContextInjection)); + } + + #[test] + fn test_message_legacy_msgpack_deser_defaults_to_none() { + // A msgpack-encoded message persisted before the `source` field + // existed must still deserialize cleanly. We synthesise the + // pre-`source` shape by encoding a plain `Message::user` and + // confirming the round-trip yields `source = None`. + let msg = Message::user("legacy"); + let bytes = rmp_serde::to_vec_named(&msg).expect("encode"); + let restored: Message = rmp_serde::from_slice(&bytes).expect("decode"); + assert_eq!(restored.source, None); + } + #[test] fn test_user_with_blocks() { let blocks = vec![ From 53a1a643e9737c731f317838dc80b15377f80f42 Mon Sep 17 00:00:00 2001 From: Philippe Branchu Date: Wed, 3 Jun 2026 04:51:29 +0000 Subject: [PATCH 02/12] fix(deps): upgrade lettre to 0.11.22 (RUSTSEC-2026-0141) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cargo update bumps lettre from 0.11.21 to 0.11.22 to clear RUSTSEC-2026-0141. Pulls in transitive dependency updates as a side effect (mostly windows-sys/socket2 version consolidation) — no API surface change in our own code. Co-Authored-By: Claude Sonnet 4.6 --- Cargo.lock | 34 +++++++++++++++++----------------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 9903cae3b9..5c79ecbfc4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -139,7 +139,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -150,7 +150,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -893,7 +893,7 @@ version = "3.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -1555,7 +1555,7 @@ dependencies = [ "libc", "option-ext", "redox_users 0.5.2", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -1805,7 +1805,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -2743,7 +2743,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.6.3", + "socket2 0.5.10", "tokio", "tower-service", "tracing", @@ -3255,9 +3255,9 @@ checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" [[package]] name = "lettre" -version = "0.11.21" +version = "0.11.22" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dabda5859ee7c06b995b9d1165aa52c39110e079ef609db97178d86aeb051fa7" +checksum = "0da65617f6cb926332d039cb578aad56178da86e128db6a1b09f4c94fa5b3349" dependencies = [ "async-trait", "base64 0.22.1", @@ -3739,7 +3739,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -5106,7 +5106,7 @@ dependencies = [ "quinn-udp", "rustc-hash", "rustls", - "socket2 0.6.3", + "socket2 0.5.10", "thiserror 2.0.18", "tokio", "tracing", @@ -5144,7 +5144,7 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2 0.6.3", + "socket2 0.5.10", "tracing", "windows-sys 0.60.2", ] @@ -5710,7 +5710,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.12.1", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -5769,7 +5769,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -6315,7 +6315,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -7051,7 +7051,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix 1.1.4", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -7623,7 +7623,7 @@ checksum = "f2f6fb2847f6742cd76af783a2a2c49e9375d0a111c7bef6f71cd9e738c72d6e" dependencies = [ "memoffset", "tempfile", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -8496,7 +8496,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] From 57cee990d750dff0abb4f90bee7736dd5e456b66 Mon Sep 17 00:00:00 2001 From: Philippe Branchu Date: Wed, 3 Jun 2026 06:02:56 +0000 Subject: [PATCH 03/12] feat(memory): structured memory storage + per-agent opt-in gate + control API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the storage layer, opt-in gate, and management/control HTTP API for structured memory. No producer (extraction + dreamer) and no UI changes — those land in follow-up PRs. Default agents see zero behavior change: storage tables are created on migration but only populated by agents that opt in via `[memory] system = "structured"` in their manifest. Per-agent opt-in: - `MemorySystem` enum (`Summarization` default / `Structured`) on `AgentManifest`, with `MemoryConfig` wrapper carrying `skip_serializing_if = is_default` so manifests that don't opt in round-trip clean TOML. - Same `[memory]` field on `HandAgentConfig`; `activate_hand` copies through to the spawned manifest. - `MemoryConfig::is_structured()` is the single gate consulted by every call site that touches structured memory. Schema migration v10 (PR 1 added v9) consolidates the three structured- memory storage tables with denormalized columns from the start: - `session_extractions(user_id, agent_id, ...)` — audit attribution survives session deletes; idx on `(user_id, created_at)` for the audit endpoint. - `user_memory_topics` with `expires_at` + `embedding` columns. - `user_agent_memory_topics` keyed by `(user_id, agent_id, topic)`. Storage modules: - `user_memory.rs` / `user_agent_memory.rs` — CRUD, embedding store, prune. - `SessionExtraction` + `SessionExtractionStore` in `session.rs`. - `MemorySubstrate::wipe_user(user_id) -> WipeUserCounts` wraps the three bucket DELETEs in a single SQLite transaction so a partial failure rolls back rather than leaving a user half-wiped. - `MemorySubstrate::list_user_extraction_audit` joins with `sessions` purely to surface the `session_deleted` flag — attribution comes from the denormalized `user_id` column. Kernel prompt-build gate: - `build_user_memory_context(memory, user_id, agent_id, memory_cfg)` returns `None` immediately when `memory_cfg.is_structured() == false` (no SQLite roundtrip for default agents). Called from both prompt-build sites in `kernel.rs`. `PromptContext::user_memory_context` carries the value through to a new "What I Remember About You" section that is skipped entirely for subagents and for empty indexes. Control API (`/api/users/*`): - `GET /api/users` — list users + default - `GET /api/users/{user_id}/memory` — list topics - `GET /api/users/{user_id}/memory/{topic}` — topic content - `DELETE /api/users/{user_id}/memory/{topic}` — delete one (404 if absent) - `DELETE /api/users/{user_id}/memory` — atomic wipe; returns per-bucket counts - `GET /api/users/{user_id}/agents/{agent_id}/memory` — per-agent topics - `DELETE /api/users/{user_id}/agents/{agent_id}/memory` — delete per-agent - `GET /api/users/{user_id}/memory/audit` — extraction events with `session_deleted` - `GET /api/users/{user_id}/memory/export` — JSON dump - `parse_user_id()` accepts `"default"` or any non-nil UUID; rejects the nil UUID (legacy anonymous-bucket sentinel) and the deprecated `"test"` alias with 400. - Module doc + per-handler `AUTHORIZATION:` comments call out the single-tenant RBAC limitation (API-key holder == full memory admin). PATCH /api/agents/:id/config gains `memory_system: Option`, validated against `MemorySystem`'s serde and persisted to disk via `Registry::update_memory_config` + `persist_manifest_to_disk`. GET /api/agents/:id surfaces both a flat `memory_system` string and the nested `manifest.memory.system` shape so dashboards can read either form. Tests (98 new): - v10 migration creates all three tables + indexes and upgrades cleanly from a v9 baseline without touching pre-existing rows. - `wipe_user`: per-bucket counts, scope (user A's wipe leaves user B untouched), idempotent zero-count run. - `MemorySystem` defaults to `Summarization`; `[memory] system = "..."` parses both variants; TOML round-trip skips `[memory]` for the default case and re-emits it when opted in. - `build_user_memory_context` returns `None` for default agents (even with seeded topics), returns the formatted block when opted in with topics, returns `None` when opted in with an empty index. - `parse_user_id`: accepts `default` and any UUID, rejects nil UUID, `"test"`, garbage, and empty string. - Audit endpoint preserves attribution after session delete and flips the `session_deleted` flag to true. --- crates/openfang-api/src/routes.rs | 666 ++++++++++++++++++ crates/openfang-api/src/server.rs | 26 + crates/openfang-hands/src/lib.rs | 11 + crates/openfang-kernel/src/heartbeat.rs | 1 + crates/openfang-kernel/src/kernel.rs | 180 +++++ crates/openfang-kernel/src/registry.rs | 20 + crates/openfang-kernel/src/wizard.rs | 1 + crates/openfang-memory/src/lib.rs | 4 +- crates/openfang-memory/src/migration.rs | 225 +++++- crates/openfang-memory/src/session.rs | 210 ++++++ crates/openfang-memory/src/substrate.rs | 424 ++++++++++- .../openfang-memory/src/user_agent_memory.rs | 386 ++++++++++ crates/openfang-memory/src/user_memory.rs | 579 +++++++++++++++ crates/openfang-runtime/src/prompt_builder.rs | 26 + crates/openfang-types/src/agent.rs | 153 ++++ 15 files changed, 2909 insertions(+), 3 deletions(-) create mode 100644 crates/openfang-memory/src/user_agent_memory.rs create mode 100644 crates/openfang-memory/src/user_memory.rs diff --git a/crates/openfang-api/src/routes.rs b/crates/openfang-api/src/routes.rs index 7bb7beecf5..3f04fca18a 100644 --- a/crates/openfang-api/src/routes.rs +++ b/crates/openfang-api/src/routes.rs @@ -1526,6 +1526,21 @@ pub async fn get_agent( "mcp_servers": entry.manifest.mcp_servers, "mcp_servers_mode": if entry.manifest.mcp_servers.is_empty() { "all" } else { "allowlist" }, "fallback_models": entry.manifest.fallback_models, + // Per-agent memory system selection. Surfaced both as a flat + // string for dashboard widgets and nested under `manifest.memory` + // so callers that round-trip the full manifest can see the field. + "memory_system": match entry.manifest.memory.system { + openfang_types::agent::MemorySystem::Summarization => "summarization", + openfang_types::agent::MemorySystem::Structured => "structured", + }, + "manifest": { + "memory": { + "system": match entry.manifest.memory.system { + openfang_types::agent::MemorySystem::Summarization => "summarization", + openfang_types::agent::MemorySystem::Structured => "structured", + }, + }, + }, })), ) } @@ -3362,6 +3377,561 @@ pub async fn get_template(Path(name): Path) -> impl IntoResponse { } } +// --------------------------------------------------------------------------- +// User memory control endpoints +// --------------------------------------------------------------------------- +// +// Expose the storage primitives in `UserMemoryStore` and `UserAgentMemoryStore` +// so users (or the dashboard) can list, view, delete, and export their +// structured memory. All endpoints require the standard API-key auth applied +// to `/api/**` routes by `middleware::auth`. +// +// # Authorization model +// +// All user memory endpoints accept any valid API key. The endpoints do NOT +// enforce that the caller is the same as the `{user_id}` in the URL — a +// single API key effectively grants access to every user's memory in the +// substrate. +// +// This matches the current single-tenant deployment model (one operator, +// one trust boundary, the API key gates everything). Multi-tenant +// deployments must enforce per-user authorization at the middleware layer +// — typically by resolving the caller's identity from the bearer token +// and rejecting requests whose URL `{user_id}` does not match the caller. +// Until such middleware lands, treat every endpoint below as "API-key +// holder == full memory admin". +// +// Each handler also carries a one-line `AUTHORIZATION:` reminder so the +// limitation is visible at the point of use, not just in this module doc. + +/// Parse and validate a user_id path parameter. +/// +/// Accepts the literal `default` (the persistent default user) or any valid +/// non-nil UUID. The nil UUID is explicitly rejected with 400 so an +/// attacker can't poke at the legacy "anonymous bucket" via the HTTP +/// surface. The deprecated `"test"` alias is also rejected — production +/// code that needs the test bucket must reach it through +/// `kernel.resolve_user_id_internal` directly, not through this function. +fn parse_user_id( + s: &str, +) -> Result)> { + if s == "default" { + return Ok(openfang_memory::session::default_user_id()); + } + match uuid::Uuid::parse_str(s) { + Ok(u) if u.is_nil() => Err(( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({ + "error": "Invalid user ID; the nil UUID is reserved and cannot be addressed via the API" + })), + )), + Ok(u) => Ok(openfang_types::agent::UserId(u)), + Err(_) => Err(( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({ + "error": "Invalid user ID; expected \"default\" or a valid UUID" + })), + )), + } +} + +/// GET /api/users — List all configured users (id, name, role). +/// +/// Includes any `[[users]]` blocks plus the persistent default user (even +/// when no block claims it). +/// +/// AUTHORIZATION: any valid API key. See module doc. +pub async fn list_users(State(state): State>) -> impl IntoResponse { + let mut users: Vec = state + .kernel + .auth + .list_users() + .into_iter() + .map(|u| { + serde_json::json!({ + "id": u.id.0.to_string(), + "name": u.name, + "role": u.role.to_string(), + }) + }) + .collect(); + + // Always surface the persistent default user, even when no [[users]] block claims it. + let default_id = state.kernel.default_user_id(); + let already_listed = users + .iter() + .any(|u| u.get("id").and_then(|v| v.as_str()) == Some(default_id.0.to_string().as_str())); + if !already_listed { + users.push(serde_json::json!({ + "id": default_id.0.to_string(), + "name": "Default User", + "role": "user", + "is_default": true, + })); + } + + ( + StatusCode::OK, + Json(serde_json::json!({ + "ok": true, + "users": users, + "default_user_id": default_id.0.to_string(), + })), + ) +} + +/// GET /api/users/:user_id/memory — List all memory topics for a user. +/// +/// AUTHORIZATION: caller is NOT required to match `{user_id}`. See module doc. +pub async fn list_user_memory( + State(state): State>, + Path(user_id_str): Path, +) -> impl IntoResponse { + let user_id = match parse_user_id(&user_id_str) { + Ok(u) => u, + Err(resp) => return resp.into_response(), + }; + + match state.kernel.memory.user_topic_index(user_id) { + Ok(entries) => { + let topics: Vec = entries + .into_iter() + .map(|e| { + serde_json::json!({ + "topic": e.topic, + "summary": e.summary, + "updated_at": e.updated_at.to_rfc3339(), + }) + }) + .collect(); + ( + StatusCode::OK, + Json(serde_json::json!({ + "ok": true, + "user_id": user_id.0.to_string(), + "topics": topics, + })), + ) + .into_response() + } + Err(e) => { + tracing::warn!("user_topic_index({user_id_str}) failed: {e}"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": "Memory operation failed"})), + ) + .into_response() + } + } +} + +/// GET /api/users/:user_id/memory/:topic — Get full content for one topic. +/// +/// AUTHORIZATION: caller is NOT required to match `{user_id}`. See module doc. +pub async fn get_user_memory_topic( + State(state): State>, + Path((user_id_str, topic)): Path<(String, String)>, +) -> impl IntoResponse { + let user_id = match parse_user_id(&user_id_str) { + Ok(u) => u, + Err(resp) => return resp.into_response(), + }; + + match state.kernel.memory.user_topic(user_id, &topic) { + Ok(Some(t)) => ( + StatusCode::OK, + Json(serde_json::json!({ + "ok": true, + "user_id": user_id.0.to_string(), + "topic": t.topic, + "summary": t.summary, + "content": t.content, + "updated_at": t.updated_at.to_rfc3339(), + "expires_at": t.expires_at.map(|d| d.to_rfc3339()), + })), + ) + .into_response(), + Ok(None) => ( + StatusCode::NOT_FOUND, + Json(serde_json::json!({"error": "Topic not found"})), + ) + .into_response(), + Err(e) => { + tracing::warn!("get_user_topic({user_id_str}, {topic}) failed: {e}"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": "Memory operation failed"})), + ) + .into_response() + } + } +} + +/// DELETE /api/users/:user_id/memory/:topic — Delete one topic for a user. +/// +/// Returns 404 if the topic does not exist (or was already expired). +/// +/// AUTHORIZATION: caller is NOT required to match `{user_id}`. See module doc. +pub async fn delete_user_memory_topic( + State(state): State>, + Path((user_id_str, topic)): Path<(String, String)>, +) -> impl IntoResponse { + let user_id = match parse_user_id(&user_id_str) { + Ok(u) => u, + Err(resp) => return resp.into_response(), + }; + + // Distinguish 404 from a successful delete by checking existence first. + let existed = matches!(state.kernel.memory.user_topic(user_id, &topic), Ok(Some(_))); + if !existed { + return ( + StatusCode::NOT_FOUND, + Json(serde_json::json!({"error": "Topic not found"})), + ) + .into_response(); + } + + match state.kernel.memory.delete_user_topic(user_id, &topic) { + Ok(()) => ( + StatusCode::OK, + Json(serde_json::json!({ + "ok": true, + "deleted": {"user_id": user_id.0.to_string(), "topic": topic}, + })), + ) + .into_response(), + Err(e) => { + tracing::warn!("delete_user_topic({user_id_str}, {topic}) failed: {e}"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": "Memory operation failed"})), + ) + .into_response() + } + } +} + +/// DELETE /api/users/:user_id/memory — Wipe all memory for a user. +/// +/// Atomic: the three bucket DELETEs run inside a single SQLite transaction +/// via `MemorySubstrate::wipe_user`, so a partial failure rolls back +/// instead of leaving a user half-wiped. +/// +/// AUTHORIZATION: caller is NOT required to match `{user_id}`. See module doc. +pub async fn delete_all_user_memory( + State(state): State>, + Path(user_id_str): Path, +) -> impl IntoResponse { + let user_id = match parse_user_id(&user_id_str) { + Ok(u) => u, + Err(resp) => return resp.into_response(), + }; + + let counts = match state.kernel.memory.wipe_user(user_id) { + Ok(c) => c, + Err(e) => { + tracing::warn!("wipe_user({user_id_str}) failed: {e}"); + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": "Memory operation failed"})), + ) + .into_response(); + } + }; + + ( + StatusCode::OK, + Json(serde_json::json!({ + "ok": true, + "user_id": user_id.0.to_string(), + "topics_deleted": counts.topics_deleted, + "agent_topics_deleted": counts.agent_topics_deleted, + "extractions_deleted": counts.extractions_deleted, + })), + ) + .into_response() +} + +/// GET /api/users/:user_id/agents/:agent_id/memory — Per-agent topics for this user. +/// +/// AUTHORIZATION: caller is NOT required to match `{user_id}`. See module doc. +pub async fn list_user_agent_memory( + State(state): State>, + Path((user_id_str, agent_id_str)): Path<(String, String)>, +) -> impl IntoResponse { + let user_id = match parse_user_id(&user_id_str) { + Ok(u) => u, + Err(resp) => return resp.into_response(), + }; + let agent_id = match uuid::Uuid::parse_str(&agent_id_str).map(AgentId) { + Ok(a) => a, + Err(_) => { + return ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({"error": "Invalid agent_id"})), + ) + .into_response(); + } + }; + + match state + .kernel + .memory + .user_agent_topic_index(user_id, agent_id) + { + Ok(entries) => { + let topics: Vec = entries + .into_iter() + .map(|e| { + serde_json::json!({ + "topic": e.topic, + "summary": e.summary, + "updated_at": e.updated_at.to_rfc3339(), + }) + }) + .collect(); + ( + StatusCode::OK, + Json(serde_json::json!({ + "ok": true, + "user_id": user_id.0.to_string(), + "agent_id": agent_id.0.to_string(), + "topics": topics, + })), + ) + .into_response() + } + Err(e) => { + tracing::warn!("user_agent_topic_index({user_id_str}, {agent_id_str}) failed: {e}"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": "Memory operation failed"})), + ) + .into_response() + } + } +} + +/// DELETE /api/users/:user_id/agents/:agent_id/memory — Delete per-agent memory. +/// +/// AUTHORIZATION: caller is NOT required to match `{user_id}`. See module doc. +pub async fn delete_user_agent_memory( + State(state): State>, + Path((user_id_str, agent_id_str)): Path<(String, String)>, +) -> impl IntoResponse { + let user_id = match parse_user_id(&user_id_str) { + Ok(u) => u, + Err(resp) => return resp.into_response(), + }; + let agent_id = match uuid::Uuid::parse_str(&agent_id_str).map(AgentId) { + Ok(a) => a, + Err(_) => { + return ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({"error": "Invalid agent_id"})), + ) + .into_response(); + } + }; + + match state + .kernel + .memory + .delete_user_agent_memory(user_id, agent_id) + { + Ok(n) => ( + StatusCode::OK, + Json(serde_json::json!({ + "ok": true, + "deleted": { + "user_id": user_id.0.to_string(), + "agent_id": agent_id.0.to_string(), + "rows": n, + }, + })), + ) + .into_response(), + Err(e) => { + tracing::warn!("delete_user_agent_memory({user_id_str}, {agent_id_str}) failed: {e}"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": "Memory operation failed"})), + ) + .into_response() + } + } +} + +/// Query parameters for `GET /api/users/:user_id/memory/audit`. +/// +/// `limit` overrides the default page size (100); the handler caps it at 500. +#[derive(serde::Deserialize)] +pub struct MemoryAuditQuery { + pub limit: Option, +} + +/// GET /api/users/:user_id/memory/audit — Recent extraction events. +/// +/// Each entry includes the originating agent, the source session id, a +/// `session_deleted` flag (true once the original session is gone — the row +/// survives via the denormalized `user_id` column), and the agent name when +/// the agent is still registered. Default limit 100, capped at 500. +/// +/// AUTHORIZATION: caller is NOT required to match `{user_id}`. See module doc. +pub async fn get_user_memory_audit( + State(state): State>, + Path(user_id_str): Path, + Query(q): Query, +) -> impl IntoResponse { + let user_id = match parse_user_id(&user_id_str) { + Ok(u) => u, + Err(resp) => return resp.into_response(), + }; + let limit = q.limit.unwrap_or(100).min(500); + + match state + .kernel + .memory + .list_user_extraction_audit(user_id, limit) + { + Ok(rows) => { + let registry = &state.kernel.registry; + let entries: Vec = rows + .into_iter() + .map( + |(id, session_id, agent_id_str, created_at, session_deleted)| { + let agent_name = uuid::Uuid::parse_str(&agent_id_str) + .ok() + .and_then(|u| registry.get(AgentId(u))) + .map(|e| e.name.clone()); + serde_json::json!({ + "id": id, + "session_id": session_id, + "agent_id": agent_id_str, + "agent_name": agent_name, + "created_at": created_at, + "session_deleted": session_deleted, + }) + }, + ) + .collect(); + ( + StatusCode::OK, + Json(serde_json::json!({ + "ok": true, + "user_id": user_id.0.to_string(), + "entries": entries, + })), + ) + .into_response() + } + Err(e) => { + tracing::warn!("list_user_extraction_audit({user_id_str}) failed: {e}"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": "Memory operation failed"})), + ) + .into_response() + } + } +} + +/// GET /api/users/:user_id/memory/export — JSON dump of everything. +/// +/// Returns user_memory topics (full content) plus per-agent topics for every +/// registered agent this user has interacted with. Suitable for user data +/// export. +/// +/// AUTHORIZATION: caller is NOT required to match `{user_id}`. See module doc. +pub async fn export_user_memory( + State(state): State>, + Path(user_id_str): Path, +) -> impl IntoResponse { + let user_id = match parse_user_id(&user_id_str) { + Ok(u) => u, + Err(resp) => return resp.into_response(), + }; + + // 1) General topics — fetch the index, then full content for each. + let index = match state.kernel.memory.user_topic_index(user_id) { + Ok(v) => v, + Err(e) => { + tracing::warn!("export: user_topic_index({user_id_str}) failed: {e}"); + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": "Memory operation failed"})), + ) + .into_response(); + } + }; + let mut topics: Vec = Vec::new(); + for e in &index { + if let Ok(Some(t)) = state.kernel.memory.user_topic(user_id, &e.topic) { + topics.push(serde_json::json!({ + "topic": t.topic, + "summary": t.summary, + "content": t.content, + "updated_at": t.updated_at.to_rfc3339(), + "expires_at": t.expires_at.map(|d| d.to_rfc3339()), + })); + } + } + + // 2) Per-agent topics. + let mut agent_topics: Vec = Vec::new(); + for entry in state.kernel.registry.list() { + let agent_id = entry.id; + if let Ok(idx) = state + .kernel + .memory + .user_agent_topic_index(user_id, agent_id) + { + if idx.is_empty() { + continue; + } + let mut full: Vec = Vec::new(); + for ie in &idx { + if let Ok(Some(t)) = state + .kernel + .memory + .user_agent_topic(user_id, agent_id, &ie.topic) + { + full.push(serde_json::json!({ + "topic": t.topic, + "summary": t.summary, + "content": t.content, + "updated_at": t.updated_at.to_rfc3339(), + })); + } + } + agent_topics.push(serde_json::json!({ + "agent_id": agent_id.0.to_string(), + "agent_name": entry.name.clone(), + "topics": full, + })); + } + } + + ( + [( + axum::http::header::CONTENT_DISPOSITION, + format!( + "attachment; filename=\"memory_{}_{}.json\"", + user_id.0, + chrono::Utc::now().format("%Y%m%d") + ), + )], + Json(serde_json::json!({ + "ok": true, + "user_id": user_id.0.to_string(), + "exported_at": chrono::Utc::now().to_rfc3339(), + "topics": topics, + "agents": agent_topics, + })), + ) + .into_response() +} + // --------------------------------------------------------------------------- // Memory endpoints // --------------------------------------------------------------------------- @@ -9708,6 +10278,9 @@ pub struct PatchAgentConfigRequest { pub api_key_env: Option, pub base_url: Option, pub fallback_models: Option>, + /// Per-agent memory system. Accepts `"summarization"` or `"structured"`. + /// Anything else returns 400. Omit to leave the current value unchanged. + pub memory_system: Option, } /// PATCH /api/agents/{id}/config — Hot-update agent name, description, system prompt, and identity. @@ -9925,6 +10498,41 @@ pub async fn patch_agent_config( } } + // Update per-agent memory system selection. + // + // Accepts only `"summarization"` (default) and `"structured"` (opt-in to + // the structured-memory pipeline). Round-trips through `MemorySystem`'s + // serde rename_all = "snake_case" via a tiny JSON detour so we get + // identical parse behavior as `[memory] system = "..."` in agent.toml. + if let Some(ref ms_str) = req.memory_system { + let parsed: Result = + serde_json::from_value(serde_json::Value::String(ms_str.clone())); + match parsed { + Ok(system) => { + let new_cfg = openfang_types::agent::MemoryConfig { system }; + if state + .kernel + .registry + .update_memory_config(agent_id, new_cfg) + .is_err() + { + return ( + StatusCode::NOT_FOUND, + Json(serde_json::json!({"error": "Agent not found"})), + ); + } + } + Err(_) => { + return ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({ + "error": "Invalid memory_system; expected \"summarization\" or \"structured\"" + })), + ); + } + } + } + // Persist updated manifest to database so changes survive restart if let Some(entry) = state.kernel.registry.get(agent_id) { if let Err(e) = state.kernel.memory.save_agent(&entry) { @@ -12975,3 +13583,61 @@ mod uninstall_agent_tests { ); } } + +#[cfg(test)] +mod user_id_parsing_tests { + use super::*; + use axum::http::StatusCode; + + #[test] + fn parse_user_id_accepts_default_keyword() { + // The literal "default" routes to the persistent default user. + let id = parse_user_id("default").expect("default keyword must parse"); + assert_eq!(id, openfang_memory::session::default_user_id()); + } + + #[test] + fn parse_user_id_accepts_valid_uuid() { + let u = uuid::Uuid::new_v4(); + let id = parse_user_id(&u.to_string()).expect("valid UUID must parse"); + assert_eq!(id.0, u); + } + + #[test] + fn parse_user_id_rejects_nil_uuid() { + // The nil UUID is the legacy anonymous-bucket sentinel and must not + // be addressable from the HTTP surface — otherwise any API-key + // holder could poke at the pre-bootstrap default-user data. + let nil = uuid::Uuid::nil(); + let err = parse_user_id(&nil.to_string()).expect_err("nil UUID must be rejected"); + assert_eq!(err.0, StatusCode::BAD_REQUEST); + let body = err.1 .0; + assert!( + body.get("error") + .and_then(|v| v.as_str()) + .unwrap_or_default() + .contains("nil UUID"), + "error message should call out the nil UUID rejection; got: {body}" + ); + } + + #[test] + fn parse_user_id_rejects_test_alias() { + // The deprecated "test" alias must not be exposed via the HTTP + // surface — tests use the in-process helper instead. + let err = parse_user_id("test").expect_err("\"test\" must be rejected"); + assert_eq!(err.0, StatusCode::BAD_REQUEST); + } + + #[test] + fn parse_user_id_rejects_garbage() { + let err = parse_user_id("not-a-uuid").expect_err("garbage must be rejected"); + assert_eq!(err.0, StatusCode::BAD_REQUEST); + } + + #[test] + fn parse_user_id_rejects_empty_string() { + let err = parse_user_id("").expect_err("empty string must be rejected"); + assert_eq!(err.0, StatusCode::BAD_REQUEST); + } +} diff --git a/crates/openfang-api/src/server.rs b/crates/openfang-api/src/server.rs index a1a2bc9c06..2f5ba48008 100644 --- a/crates/openfang-api/src/server.rs +++ b/crates/openfang-api/src/server.rs @@ -295,6 +295,32 @@ pub async fn build_router( "/api/uploads/{file_id}", axum::routing::get(routes::serve_upload), ) + // User memory control endpoints (see routes::list_user_memory module doc). + // The `/memory/audit` route comes BEFORE the `:topic` capture so it + // doesn't get swallowed by the catch-all topic parameter. + .route("/api/users", axum::routing::get(routes::list_users)) + .route( + "/api/users/{user_id}/memory", + axum::routing::get(routes::list_user_memory).delete(routes::delete_all_user_memory), + ) + .route( + "/api/users/{user_id}/memory/audit", + axum::routing::get(routes::get_user_memory_audit), + ) + .route( + "/api/users/{user_id}/memory/export", + axum::routing::get(routes::export_user_memory), + ) + .route( + "/api/users/{user_id}/memory/{topic}", + axum::routing::get(routes::get_user_memory_topic) + .delete(routes::delete_user_memory_topic), + ) + .route( + "/api/users/{user_id}/agents/{agent_id}/memory", + axum::routing::get(routes::list_user_agent_memory) + .delete(routes::delete_user_agent_memory), + ) // Channel endpoints .route("/api/channels", axum::routing::get(routes::list_channels)) .route( diff --git a/crates/openfang-hands/src/lib.rs b/crates/openfang-hands/src/lib.rs index 1bdb6783c3..220e12468e 100644 --- a/crates/openfang-hands/src/lib.rs +++ b/crates/openfang-hands/src/lib.rs @@ -298,6 +298,17 @@ pub struct HandAgentConfig { /// making long LLM calls. Omit to use the kernel default. #[serde(default)] pub heartbeat_interval_secs: Option, + /// Per-hand memory system selection. Omitted in HAND.toml → defaults + /// to `MemorySystem::Summarization` (current OpenFang behavior). + /// + /// Set to `structured` to enable the structured memory pipeline + /// (extractions + dream consolidation) for this hand's spawned agent. + /// `activate_hand` copies this through to the spawned `AgentManifest`. + #[serde( + default, + skip_serializing_if = "openfang_types::agent::MemoryConfig::is_default" + )] + pub memory: openfang_types::agent::MemoryConfig, } fn default_module() -> String { diff --git a/crates/openfang-kernel/src/heartbeat.rs b/crates/openfang-kernel/src/heartbeat.rs index 75554b1da6..ab5e66b8a3 100644 --- a/crates/openfang-kernel/src/heartbeat.rs +++ b/crates/openfang-kernel/src/heartbeat.rs @@ -346,6 +346,7 @@ mod tests { tool_blocklist: vec![], cache_context: false, max_history_messages: None, + memory: openfang_types::agent::MemoryConfig::default(), }, state, mode: AgentMode::default(), diff --git a/crates/openfang-kernel/src/kernel.rs b/crates/openfang-kernel/src/kernel.rs index f520edeb88..d21104149d 100644 --- a/crates/openfang-kernel/src/kernel.rs +++ b/crates/openfang-kernel/src/kernel.rs @@ -526,6 +526,71 @@ fn read_identity_file(state_dir: &Path, filename: &str) -> Option { } } +/// Build the user-memory context block for injection into the system prompt. +/// +/// Combines general user topics and per-agent user topics into a concise +/// index with a human-readable "age" for staleness reasoning. The block is +/// rendered later by `prompt_builder` under the `## What I Remember About +/// You` heading. +/// +/// **Opt-in gate:** returns `None` immediately (no DB roundtrip) when the +/// agent has not opted in to structured memory +/// (`MemoryConfig::is_structured() == false`). The default +/// `MemorySystem::Summarization` writes nothing to the user-memory stores, +/// so the result would always be empty anyway — exiting early avoids the +/// SQLite read on every prompt build for upstream-shape agents. +fn build_user_memory_context( + memory: &openfang_memory::MemorySubstrate, + user_id: openfang_types::agent::UserId, + agent_id: AgentId, + memory_cfg: &openfang_types::agent::MemoryConfig, +) -> Option { + if !memory_cfg.is_structured() { + return None; + } + let user_topics = memory.user_topic_index(user_id).unwrap_or_default(); + let agent_topics = memory + .user_agent_topic_index(user_id, agent_id) + .unwrap_or_default(); + + if user_topics.is_empty() && agent_topics.is_empty() { + return None; + } + + let mut lines = Vec::new(); + for t in &user_topics { + let age = format_memory_age(t.updated_at); + lines.push(format!("- **{}** — {} ({})", t.topic, t.summary, age)); + } + for t in &agent_topics { + let age = format_memory_age(t.updated_at); + lines.push(format!( + "- **{}** [with me] — {} ({})", + t.topic, t.summary, age + )); + } + Some(lines.join("\n")) +} + +/// Format a UTC timestamp as a human-readable age string +/// (e.g. "3 days ago", "just now"). Used by `build_user_memory_context`. +fn format_memory_age(updated_at: chrono::DateTime) -> String { + let secs = (chrono::Utc::now() - updated_at).num_seconds(); + if secs < 60 { + "just now".to_string() + } else if secs < 3600 { + format!("{} min ago", secs / 60) + } else if secs < 86400 { + format!("{} hr ago", secs / 3600) + } else if secs < 86400 * 7 { + format!("{} days ago", secs / 86400) + } else if secs < 86400 * 30 { + format!("{} weeks ago", secs / (86400 * 7)) + } else { + format!("{} months ago", secs / (86400 * 30)) + } +} + /// Get the system hostname as a String. fn gethostname() -> Option { #[cfg(unix)] @@ -2302,6 +2367,15 @@ impl OpenFangKernel { context_md: manifest.workspace.as_ref().and_then(|w| { openfang_runtime::agent_context::load_context_md(w, manifest.cache_context) }), + // Gated on `manifest.memory.is_structured()` — returns None + // immediately for default-Summarization agents so the SQLite + // round-trip is skipped on the hot path. + user_memory_context: build_user_memory_context( + &self.memory, + session.user_id, + agent_id, + &manifest.memory, + ), }; manifest.model.system_prompt = openfang_runtime::prompt_builder::build_system_prompt(&prompt_ctx); @@ -2885,6 +2959,14 @@ impl OpenFangKernel { context_md: manifest.workspace.as_ref().and_then(|w| { openfang_runtime::agent_context::load_context_md(w, manifest.cache_context) }), + // Gated on `manifest.memory.is_structured()` — see the helper + // for the early-return rationale. + user_memory_context: build_user_memory_context( + &self.memory, + session.user_id, + agent_id, + &manifest.memory, + ), }; manifest.model.system_prompt = openfang_runtime::prompt_builder::build_system_prompt(&prompt_ctx); @@ -3890,6 +3972,8 @@ impl OpenFangKernel { } else { None }, + // Propagate hand's per-agent memory opt-in to the spawned manifest. + memory: def.agent.memory.clone(), ..Default::default() }; @@ -8190,6 +8274,7 @@ mod tests { tool_blocklist: vec![], cache_context: false, max_history_messages: None, + memory: openfang_types::agent::MemoryConfig::default(), }; manifest.capabilities.tools = vec!["file_read".to_string(), "web_fetch".to_string()]; manifest.capabilities.agent_spawn = true; @@ -8234,6 +8319,7 @@ mod tests { tool_blocklist: vec![], cache_context: false, max_history_messages: None, + memory: openfang_types::agent::MemoryConfig::default(), }; let mut disk = entry.clone(); disk.description = "new".to_string(); @@ -8285,6 +8371,7 @@ mod tests { tool_blocklist: vec![], cache_context: false, max_history_messages: None, + memory: openfang_types::agent::MemoryConfig::default(), }; let mut disk = entry.clone(); disk.workspace = Some(std::path::PathBuf::from("/new")); @@ -8342,6 +8429,7 @@ mod tests { tool_blocklist: vec![], cache_context: false, max_history_messages: None, + memory: openfang_types::agent::MemoryConfig::default(), }; // Current kernel config now says mode = Full. @@ -8456,6 +8544,7 @@ mod tests { tool_blocklist: vec![], cache_context: false, max_history_messages: None, + memory: openfang_types::agent::MemoryConfig::default(), } } @@ -9859,4 +9948,95 @@ system_prompt = "You are a test agent." kernel.shutdown(); } + + // ----------------------------------------------------------------- + // build_user_memory_context — the opt-in gate that decides whether + // structured memory bleeds into the system prompt. + // ----------------------------------------------------------------- + + /// Default agents (no `[memory]` block, or `system = "summarization"`) + /// MUST return `None` from the helper without touching SQLite. This is the + /// core upstream-compatibility guarantee of the structured-memory feature. + #[test] + fn test_build_user_memory_context_returns_none_for_default_agents() { + let tmp = tempfile::tempdir().unwrap(); + let kernel = minimal_kernel(&tmp); + let agent_id = openfang_types::agent::AgentId::new(); + let user_id = openfang_types::agent::UserId::new(); + let default_cfg = openfang_types::agent::MemoryConfig::default(); + + // Seed a topic — the helper must STILL return None because the gate + // is on the memory config, not on the contents of the store. + kernel + .memory + .upsert_user_topic(&openfang_memory::user_memory::MemoryTopic { + user_id, + topic: "prefs".into(), + summary: "summary".into(), + content: "content".into(), + updated_at: chrono::Utc::now(), + expires_at: None, + }) + .unwrap(); + + let result = build_user_memory_context(&kernel.memory, user_id, agent_id, &default_cfg); + assert!( + result.is_none(), + "default-Summarization agents must skip the user-memory block" + ); + + kernel.shutdown(); + } + + /// Opted-in agents WITH topics return the formatted index. + #[test] + fn test_build_user_memory_context_populated_for_opted_in_agents() { + let tmp = tempfile::tempdir().unwrap(); + let kernel = minimal_kernel(&tmp); + let agent_id = openfang_types::agent::AgentId::new(); + let user_id = openfang_types::agent::UserId::new(); + let structured_cfg = openfang_types::agent::MemoryConfig { + system: openfang_types::agent::MemorySystem::Structured, + }; + + kernel + .memory + .upsert_user_topic(&openfang_memory::user_memory::MemoryTopic { + user_id, + topic: "prefs".into(), + summary: "likes tea".into(), + content: "details".into(), + updated_at: chrono::Utc::now(), + expires_at: None, + }) + .unwrap(); + + let result = + build_user_memory_context(&kernel.memory, user_id, agent_id, &structured_cfg).unwrap(); + assert!(result.contains("prefs")); + assert!(result.contains("likes tea")); + + kernel.shutdown(); + } + + /// Opted-in agent with no stored topics returns `None` (empty index + /// shouldn't add a noisy empty section to the prompt). + #[test] + fn test_build_user_memory_context_none_when_no_topics() { + let tmp = tempfile::tempdir().unwrap(); + let kernel = minimal_kernel(&tmp); + let agent_id = openfang_types::agent::AgentId::new(); + let user_id = openfang_types::agent::UserId::new(); + let structured_cfg = openfang_types::agent::MemoryConfig { + system: openfang_types::agent::MemorySystem::Structured, + }; + + let result = build_user_memory_context(&kernel.memory, user_id, agent_id, &structured_cfg); + assert!( + result.is_none(), + "empty index should yield None, not an empty block" + ); + + kernel.shutdown(); + } } diff --git a/crates/openfang-kernel/src/registry.rs b/crates/openfang-kernel/src/registry.rs index 124bbfe728..1a0b3a707c 100644 --- a/crates/openfang-kernel/src/registry.rs +++ b/crates/openfang-kernel/src/registry.rs @@ -327,6 +327,25 @@ impl AgentRegistry { Ok(()) } + /// Update an agent's memory configuration (opt-in to/out of structured memory). + /// + /// Hot-swappable: the next prompt build picks up the new gate. Persisting + /// to disk is the caller's responsibility (the control-API handler calls + /// `persist_manifest_to_disk` right after). + pub fn update_memory_config( + &self, + id: AgentId, + new_memory: openfang_types::agent::MemoryConfig, + ) -> OpenFangResult<()> { + let mut entry = self + .agents + .get_mut(&id) + .ok_or_else(|| OpenFangError::AgentNotFound(id.to_string()))?; + entry.manifest.memory = new_memory; + entry.last_active = chrono::Utc::now(); + Ok(()) + } + /// Update an agent's resource quota (budget limits). pub fn update_resources( &self, @@ -415,6 +434,7 @@ mod tests { tool_blocklist: vec![], cache_context: false, max_history_messages: None, + memory: openfang_types::agent::MemoryConfig::default(), }, state: AgentState::Created, mode: AgentMode::default(), diff --git a/crates/openfang-kernel/src/wizard.rs b/crates/openfang-kernel/src/wizard.rs index 4338830af2..7f18147fe1 100644 --- a/crates/openfang-kernel/src/wizard.rs +++ b/crates/openfang-kernel/src/wizard.rs @@ -185,6 +185,7 @@ impl SetupWizard { tool_blocklist: vec![], cache_context: false, max_history_messages: None, + memory: openfang_types::agent::MemoryConfig::default(), }; let skills_to_install: Vec = intent diff --git a/crates/openfang-memory/src/lib.rs b/crates/openfang-memory/src/lib.rs index bf16737a02..b8fbb2340f 100644 --- a/crates/openfang-memory/src/lib.rs +++ b/crates/openfang-memory/src/lib.rs @@ -16,6 +16,8 @@ pub mod semantic; pub mod session; pub mod structured; pub mod usage; +pub mod user_agent_memory; +pub mod user_memory; mod substrate; -pub use substrate::MemorySubstrate; +pub use substrate::{ExtractionAuditRow, MemorySubstrate, WipeUserCounts}; diff --git a/crates/openfang-memory/src/migration.rs b/crates/openfang-memory/src/migration.rs index 88c64f77a3..5010ec07c8 100644 --- a/crates/openfang-memory/src/migration.rs +++ b/crates/openfang-memory/src/migration.rs @@ -5,7 +5,7 @@ use rusqlite::Connection; /// Current schema version. -const SCHEMA_VERSION: u32 = 9; +const SCHEMA_VERSION: u32 = 10; /// Run all migrations to bring the database up to date. pub fn run_migrations(conn: &Connection) -> Result<(), rusqlite::Error> { @@ -47,6 +47,10 @@ pub fn run_migrations(conn: &Connection) -> Result<(), rusqlite::Error> { migrate_v9(conn)?; } + if current_version < 10 { + migrate_v10(conn)?; + } + set_schema_version(conn, SCHEMA_VERSION)?; Ok(()) } @@ -366,6 +370,90 @@ fn migrate_v9(conn: &Connection) -> Result<(), rusqlite::Error> { Ok(()) } +/// Version 10: Structured-memory storage tables. +/// +/// Adds three tables — `session_extractions`, `user_memory_topics`, +/// `user_agent_memory_topics` — that back the opt-in structured memory +/// system. None of these tables are populated by default agents: the +/// producer (`extract_structured` / dreamer) is gated on +/// `MemoryConfig::is_structured()` and lands in a later PR. The tables are +/// created unconditionally so the storage and control-API surface is +/// available the moment an agent opts in. +/// +/// Schema choices: +/// +/// - `session_extractions` carries denormalized `user_id` and `agent_id` +/// columns from the start. The audit endpoint filters by owning user +/// and surfaces a `session_deleted` flag derived from a LEFT JOIN +/// against `sessions`; keeping the ids on the row means audit +/// attribution survives a session delete instead of orphaning the row. +/// - `user_memory_topics` ships with `expires_at` (optional ISO timestamp; +/// expired rows are pruned at read time) and `embedding` (BLOB of packed +/// little-endian f32 values for optional cosine-similarity retrieval). +/// The producer populates `embedding` asynchronously, so the column +/// stays nullable. +/// - `user_agent_memory_topics` is the per-(user, agent) sibling of +/// `user_memory_topics`. Topics here are scoped to "how this user +/// interacts with this specific agent" and do not bleed across agents. +/// +/// Indexes: +/// - `idx_extractions_session` for the per-session loader. +/// - `idx_extractions_user_created` for the audit endpoint's +/// `(user_id, created_at)` ordering. +/// - `idx_user_memory_user` for the per-user topic index. +/// - `idx_user_agent_memory` for the per-(user, agent) topic index. +fn migrate_v10(conn: &Connection) -> Result<(), rusqlite::Error> { + conn.execute_batch( + " + CREATE TABLE IF NOT EXISTS session_extractions ( + id TEXT PRIMARY KEY, + session_id TEXT NOT NULL, + user_id TEXT NOT NULL DEFAULT '00000000-0000-0000-0000-000000000000', + agent_id TEXT NOT NULL DEFAULT '00000000-0000-0000-0000-000000000000', + facts TEXT NOT NULL DEFAULT '[]', + preferences TEXT NOT NULL DEFAULT '[]', + decisions TEXT NOT NULL DEFAULT '[]', + tasks TEXT NOT NULL DEFAULT '[]', + open_items TEXT NOT NULL DEFAULT '[]', + created_at TEXT NOT NULL + ); + CREATE INDEX IF NOT EXISTS idx_extractions_session + ON session_extractions(session_id); + CREATE INDEX IF NOT EXISTS idx_extractions_user_created + ON session_extractions(user_id, created_at); + + CREATE TABLE IF NOT EXISTS user_memory_topics ( + user_id TEXT NOT NULL, + topic TEXT NOT NULL, + summary TEXT NOT NULL DEFAULT '', + content TEXT NOT NULL DEFAULT '', + updated_at TEXT NOT NULL, + expires_at TEXT DEFAULT NULL, + embedding BLOB DEFAULT NULL, + PRIMARY KEY (user_id, topic) + ); + CREATE INDEX IF NOT EXISTS idx_user_memory_user + ON user_memory_topics(user_id); + + CREATE TABLE IF NOT EXISTS user_agent_memory_topics ( + user_id TEXT NOT NULL, + agent_id TEXT NOT NULL, + topic TEXT NOT NULL, + summary TEXT NOT NULL DEFAULT '', + content TEXT NOT NULL DEFAULT '', + updated_at TEXT NOT NULL, + PRIMARY KEY (user_id, agent_id, topic) + ); + CREATE INDEX IF NOT EXISTS idx_user_agent_memory + ON user_agent_memory_topics(user_id, agent_id); + + INSERT OR IGNORE INTO migrations (version, applied_at, description) + VALUES (10, datetime('now'), 'Structured memory storage tables (session_extractions, user_memory_topics, user_agent_memory_topics)'); + ", + )?; + Ok(()) +} + #[cfg(test)] mod tests { use super::*; @@ -472,4 +560,139 @@ mod tests { assert_eq!(uid, "00000000-0000-0000-0000-000000000000"); assert!(parent.is_none()); } + + // ── v10: Structured-memory storage ────────────────────────────────── + + #[test] + fn test_migration_v10_creates_storage_tables() { + let conn = Connection::open_in_memory().unwrap(); + run_migrations(&conn).unwrap(); + let tables: Vec = conn + .prepare("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name") + .unwrap() + .query_map([], |row| row.get(0)) + .unwrap() + .filter_map(|r| r.ok()) + .collect(); + assert!(tables.contains(&"session_extractions".to_string())); + assert!(tables.contains(&"user_memory_topics".to_string())); + assert!(tables.contains(&"user_agent_memory_topics".to_string())); + } + + #[test] + fn test_migration_v10_session_extractions_columns_denormalized() { + // user_id and agent_id are denormalized onto the row from the start — + // no v15 backfill needed because PR 2 ships the table this way. + let conn = Connection::open_in_memory().unwrap(); + run_migrations(&conn).unwrap(); + assert!(column_exists(&conn, "session_extractions", "user_id")); + assert!(column_exists(&conn, "session_extractions", "agent_id")); + assert!(column_exists(&conn, "session_extractions", "session_id")); + assert!(column_exists(&conn, "session_extractions", "facts")); + assert!(column_exists(&conn, "session_extractions", "preferences")); + assert!(column_exists(&conn, "session_extractions", "decisions")); + assert!(column_exists(&conn, "session_extractions", "tasks")); + assert!(column_exists(&conn, "session_extractions", "open_items")); + assert!(column_exists(&conn, "session_extractions", "created_at")); + } + + #[test] + fn test_migration_v10_user_memory_topics_columns() { + let conn = Connection::open_in_memory().unwrap(); + run_migrations(&conn).unwrap(); + // expires_at and embedding are present from the start so prune-at-read + // and similarity-search code paths compile and run. + assert!(column_exists(&conn, "user_memory_topics", "user_id")); + assert!(column_exists(&conn, "user_memory_topics", "topic")); + assert!(column_exists(&conn, "user_memory_topics", "summary")); + assert!(column_exists(&conn, "user_memory_topics", "content")); + assert!(column_exists(&conn, "user_memory_topics", "updated_at")); + assert!(column_exists(&conn, "user_memory_topics", "expires_at")); + assert!(column_exists(&conn, "user_memory_topics", "embedding")); + } + + #[test] + fn test_migration_v10_user_agent_memory_topics_columns() { + let conn = Connection::open_in_memory().unwrap(); + run_migrations(&conn).unwrap(); + assert!(column_exists(&conn, "user_agent_memory_topics", "user_id")); + assert!(column_exists(&conn, "user_agent_memory_topics", "agent_id")); + assert!(column_exists(&conn, "user_agent_memory_topics", "topic")); + assert!(column_exists(&conn, "user_agent_memory_topics", "summary")); + assert!(column_exists(&conn, "user_agent_memory_topics", "content")); + assert!(column_exists( + &conn, + "user_agent_memory_topics", + "updated_at" + )); + } + + #[test] + fn test_migration_v10_indexes_present() { + let conn = Connection::open_in_memory().unwrap(); + run_migrations(&conn).unwrap(); + let names: Vec = conn + .prepare("SELECT name FROM sqlite_master WHERE type='index'") + .unwrap() + .query_map([], |row| row.get(0)) + .unwrap() + .filter_map(|r| r.ok()) + .collect(); + assert!(names.contains(&"idx_extractions_session".to_string())); + assert!(names.contains(&"idx_extractions_user_created".to_string())); + assert!(names.contains(&"idx_user_memory_user".to_string())); + assert!(names.contains(&"idx_user_agent_memory".to_string())); + } + + /// A v9 database (PR 1 tip) must upgrade cleanly to v10: the three new + /// tables appear without touching pre-existing rows. + #[test] + fn test_migration_v9_to_v10_upgrade_preserves_existing_rows() { + let conn = Connection::open_in_memory().unwrap(); + + // Stop at v9 (PR 1 baseline). + migrate_v1(&conn).unwrap(); + migrate_v2(&conn).unwrap(); + migrate_v3(&conn).unwrap(); + migrate_v4(&conn).unwrap(); + migrate_v5(&conn).unwrap(); + migrate_v6(&conn).unwrap(); + migrate_v7(&conn).unwrap(); + migrate_v8(&conn).unwrap(); + migrate_v9(&conn).unwrap(); + set_schema_version(&conn, 9).unwrap(); + + // Insert a v9 session row — nothing here should be touched by v10. + conn.execute( + "INSERT INTO sessions (id, agent_id, messages, context_window_tokens, user_id, created_at, updated_at) \ + VALUES ('sess-v9', 'agent-v9', X'', 0, '11111111-1111-1111-1111-111111111111', datetime('now'), datetime('now'))", + [], + ) + .unwrap(); + + // Apply v10. + run_migrations(&conn).unwrap(); + + // New tables exist. + let tables: Vec = conn + .prepare("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name") + .unwrap() + .query_map([], |row| row.get(0)) + .unwrap() + .filter_map(|r| r.ok()) + .collect(); + assert!(tables.contains(&"session_extractions".to_string())); + assert!(tables.contains(&"user_memory_topics".to_string())); + assert!(tables.contains(&"user_agent_memory_topics".to_string())); + + // Pre-existing session row untouched. + let uid: String = conn + .query_row( + "SELECT user_id FROM sessions WHERE id = 'sess-v9'", + [], + |row| row.get(0), + ) + .unwrap(); + assert_eq!(uid, "11111111-1111-1111-1111-111111111111"); + } } diff --git a/crates/openfang-memory/src/session.rs b/crates/openfang-memory/src/session.rs index 1176a58e50..c69c9ffb0c 100644 --- a/crates/openfang-memory/src/session.rs +++ b/crates/openfang-memory/src/session.rs @@ -214,6 +214,57 @@ impl SessionStore { Ok(()) } + /// Get all direct child session IDs for a parent session. + /// + /// Uses the `parent_session_id` column added in v9 (PR 1). The producer + /// that actually forks sessions (hand sessions linked to callers) lands + /// in a later PR — until then this returns empty for every input. + pub fn get_child_sessions( + &self, + parent_session_id: SessionId, + ) -> OpenFangResult> { + let conn = self + .conn + .lock() + .map_err(|e| OpenFangError::Internal(e.to_string()))?; + let mut stmt = conn + .prepare("SELECT id FROM sessions WHERE parent_session_id = ?1") + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + + let rows = stmt + .query_map(rusqlite::params![parent_session_id.0.to_string()], |row| { + row.get::<_, String>(0) + }) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + + let mut ids = Vec::new(); + for row in rows { + let id_str = row.map_err(|e| OpenFangError::Memory(e.to_string()))?; + let session_id = uuid::Uuid::parse_str(&id_str) + .map(SessionId) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + ids.push(session_id); + } + Ok(ids) + } + + /// Delete a session and all its descendants recursively. + /// + /// Returns the total number of sessions deleted (including the root). + /// No forks exist today so this is functionally equivalent to + /// `delete_session`; once the fork producer lands, every descendant is + /// pulled in via the `parent_session_id` link. + pub fn delete_session_tree(&self, root_session_id: SessionId) -> OpenFangResult { + let children = self.get_child_sessions(root_session_id)?; + let mut count = 0; + for child_id in children { + count += self.delete_session_tree(child_id)?; + } + self.delete_session(root_session_id)?; + count += 1; + Ok(count) + } + /// Delete all sessions belonging to an agent. pub fn delete_agent_sessions(&self, agent_id: AgentId) -> OpenFangResult<()> { let conn = self @@ -750,6 +801,165 @@ impl SessionStore { } } +/// A structured extraction pass result stored mid-session during compaction. +/// +/// Populated by the (not-yet-landed) producer in a later PR. This PR provides +/// the type and the storage backend so the surrounding plumbing — control +/// API, wipe, audit — can be wired up independently of the producer. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, Default)] +pub struct SessionExtraction { + pub facts: Vec, + pub preferences: Vec, + pub decisions: Vec, + pub tasks: Vec, + pub open_items: Vec, +} + +/// Store for session extractions (compaction outputs). +/// +/// Each `append` denormalizes the session's owning `user_id` and `agent_id` +/// onto the extraction row so the audit endpoint can attribute the event +/// even after the originating session is deleted. The audit query joins +/// against `sessions` purely to surface the `session_deleted` flag, not for +/// attribution. +#[derive(Clone)] +pub struct SessionExtractionStore { + conn: Arc>, +} + +impl SessionExtractionStore { + pub fn new(conn: Arc>) -> Self { + Self { conn } + } + + /// Append a new extraction record for a session. + /// + /// Looks up the owning `user_id`/`agent_id` from the `sessions` table and + /// denormalizes them onto the extraction row so the audit endpoint can + /// attribute the event correctly even after the session is deleted. If + /// the session is missing (shouldn't happen in normal flow — extractions + /// are written by code that just loaded the session) we fall back to the + /// nil UUID. + pub fn append( + &self, + session_id: SessionId, + extraction: &SessionExtraction, + ) -> OpenFangResult<()> { + let conn = self + .conn + .lock() + .map_err(|e| OpenFangError::Internal(e.to_string()))?; + let id = uuid::Uuid::new_v4().to_string(); + let now = Utc::now().to_rfc3339(); + let facts = serde_json::to_string(&extraction.facts) + .map_err(|e| OpenFangError::Serialization(e.to_string()))?; + let preferences = serde_json::to_string(&extraction.preferences) + .map_err(|e| OpenFangError::Serialization(e.to_string()))?; + let decisions = serde_json::to_string(&extraction.decisions) + .map_err(|e| OpenFangError::Serialization(e.to_string()))?; + let tasks = serde_json::to_string(&extraction.tasks) + .map_err(|e| OpenFangError::Serialization(e.to_string()))?; + let open_items = serde_json::to_string(&extraction.open_items) + .map_err(|e| OpenFangError::Serialization(e.to_string()))?; + + // Resolve owning user/agent for the audit-friendly denormalized columns. + let nil = uuid::Uuid::nil().to_string(); + let (user_id, agent_id): (String, String) = conn + .query_row( + "SELECT user_id, agent_id FROM sessions WHERE id = ?1", + rusqlite::params![session_id.0.to_string()], + |row| { + let u: Option = row.get(0).ok(); + let a: Option = row.get(1).ok(); + Ok(( + u.unwrap_or_else(|| nil.clone()), + a.unwrap_or_else(|| nil.clone()), + )) + }, + ) + .unwrap_or_else(|_| (nil.clone(), nil.clone())); + + conn.execute( + "INSERT INTO session_extractions \ + (id, session_id, facts, preferences, decisions, tasks, open_items, created_at, user_id, agent_id) \ + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)", + rusqlite::params![ + id, + session_id.0.to_string(), + facts, + preferences, + decisions, + tasks, + open_items, + now, + user_id, + agent_id, + ], + ) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + Ok(()) + } + + /// Load all extractions for a session, ordered by created_at. + pub fn load_all(&self, session_id: SessionId) -> OpenFangResult> { + let conn = self + .conn + .lock() + .map_err(|e| OpenFangError::Internal(e.to_string()))?; + let mut stmt = conn + .prepare( + "SELECT facts, preferences, decisions, tasks, open_items \ + FROM session_extractions WHERE session_id = ?1 ORDER BY created_at ASC", + ) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + + let rows = stmt + .query_map(rusqlite::params![session_id.0.to_string()], |row| { + let facts: String = row.get(0)?; + let preferences: String = row.get(1)?; + let decisions: String = row.get(2)?; + let tasks: String = row.get(3)?; + let open_items: String = row.get(4)?; + Ok((facts, preferences, decisions, tasks, open_items)) + }) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + + let mut extractions = Vec::new(); + for row in rows { + let (facts, preferences, decisions, tasks, open_items) = + row.map_err(|e| OpenFangError::Memory(e.to_string()))?; + let extraction = SessionExtraction { + facts: serde_json::from_str(&facts) + .map_err(|e| OpenFangError::Serialization(e.to_string()))?, + preferences: serde_json::from_str(&preferences) + .map_err(|e| OpenFangError::Serialization(e.to_string()))?, + decisions: serde_json::from_str(&decisions) + .map_err(|e| OpenFangError::Serialization(e.to_string()))?, + tasks: serde_json::from_str(&tasks) + .map_err(|e| OpenFangError::Serialization(e.to_string()))?, + open_items: serde_json::from_str(&open_items) + .map_err(|e| OpenFangError::Serialization(e.to_string()))?, + }; + extractions.push(extraction); + } + Ok(extractions) + } + + /// Delete all extractions for a session. + pub fn delete_for_session(&self, session_id: SessionId) -> OpenFangResult<()> { + let conn = self + .conn + .lock() + .map_err(|e| OpenFangError::Internal(e.to_string()))?; + conn.execute( + "DELETE FROM session_extractions WHERE session_id = ?1", + rusqlite::params![session_id.0.to_string()], + ) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + Ok(()) + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/openfang-memory/src/substrate.rs b/crates/openfang-memory/src/substrate.rs index 41dc160bba..b49a812249 100644 --- a/crates/openfang-memory/src/substrate.rs +++ b/crates/openfang-memory/src/substrate.rs @@ -7,9 +7,13 @@ use crate::consolidation::ConsolidationEngine; use crate::knowledge::KnowledgeStore; use crate::migration::run_migrations; use crate::semantic::SemanticStore; -use crate::session::{Session, SessionStore}; +use crate::session::{Session, SessionExtraction, SessionExtractionStore, SessionStore}; use crate::structured::StructuredStore; use crate::usage::UsageStore; +use crate::user_agent_memory::{ + UserAgentMemoryStore, UserAgentMemoryTopic, UserAgentTopicIndexEntry, +}; +use crate::user_memory::{MemoryTopic, TopicIndexEntry, UserMemoryStore}; use async_trait::async_trait; use openfang_types::agent::{AgentEntry, AgentId, SessionId, UserId}; @@ -25,6 +29,29 @@ use std::path::Path; use std::sync::{Arc, Mutex}; use tracing::{info, warn}; +/// Per-bucket row counts returned by [`MemorySubstrate::wipe_user`]. +/// +/// Mirrors the response body of `DELETE /api/users/:user_id/memory` so the +/// handler can forward the struct verbatim. All three buckets are reported +/// even when zero so the UI can render "0 deleted" rather than guessing. +#[derive(Debug, Clone, serde::Serialize)] +pub struct WipeUserCounts { + /// Rows deleted from `user_memory_topics`. + pub topics_deleted: usize, + /// Rows deleted from `user_agent_memory_topics`. + pub agent_topics_deleted: usize, + /// Rows deleted from `session_extractions` (audit log). + pub extractions_deleted: usize, +} + +/// Row shape returned by [`MemorySubstrate::list_user_extraction_audit`]. +/// +/// Fields, in order: `extraction_id`, `session_id`, `agent_id`, +/// `created_at_rfc3339`, `session_deleted`. The control-API handler maps +/// this to a JSON entry; keeping the substrate-layer return as a tuple +/// keeps the substrate JSON-free. +pub type ExtractionAuditRow = (String, String, String, String, bool); + /// The unified memory substrate. Implements the `Memory` trait by delegating /// to specialized stores backed by a shared SQLite connection. pub struct MemorySubstrate { @@ -33,6 +60,9 @@ pub struct MemorySubstrate { semantic: SemanticStore, knowledge: KnowledgeStore, sessions: SessionStore, + extractions: SessionExtractionStore, + user_memory: UserMemoryStore, + user_agent_memory: UserAgentMemoryStore, consolidation: ConsolidationEngine, usage: UsageStore, } @@ -62,6 +92,9 @@ impl MemorySubstrate { semantic, knowledge: KnowledgeStore::new(Arc::clone(&shared)), sessions: SessionStore::new(Arc::clone(&shared)), + extractions: SessionExtractionStore::new(Arc::clone(&shared)), + user_memory: UserMemoryStore::new(Arc::clone(&shared)), + user_agent_memory: UserAgentMemoryStore::new(Arc::clone(&shared)), usage: UsageStore::new(Arc::clone(&shared)), consolidation: ConsolidationEngine::new(shared, decay_rate), }) @@ -116,6 +149,9 @@ impl MemorySubstrate { semantic: SemanticStore::new(Arc::clone(&shared)), knowledge: KnowledgeStore::new(Arc::clone(&shared)), sessions: SessionStore::new(Arc::clone(&shared)), + extractions: SessionExtractionStore::new(Arc::clone(&shared)), + user_memory: UserMemoryStore::new(Arc::clone(&shared)), + user_agent_memory: UserAgentMemoryStore::new(Arc::clone(&shared)), usage: UsageStore::new(Arc::clone(&shared)), consolidation: ConsolidationEngine::new(shared, decay_rate), }) @@ -356,6 +392,262 @@ impl MemorySubstrate { Ok(()) } + // ----------------------------------------------------------------- + // Structured memory: extractions + // ----------------------------------------------------------------- + + /// Append a structured extraction for a session. + /// + /// No producer in this PR — this is the storage surface that PR 3's + /// `extract_structured` calls into. Default agents never reach this + /// path because the producer is gated on `MemoryConfig::is_structured()`. + pub fn append_extraction( + &self, + session_id: SessionId, + extraction: &SessionExtraction, + ) -> OpenFangResult<()> { + self.extractions.append(session_id, extraction) + } + + /// Load all structured extractions for a session, oldest first. + pub fn load_extractions( + &self, + session_id: SessionId, + ) -> OpenFangResult> { + self.extractions.load_all(session_id) + } + + /// Delete all extractions for a session. + pub fn delete_extractions(&self, session_id: SessionId) -> OpenFangResult<()> { + self.extractions.delete_for_session(session_id) + } + + /// Delete every extraction event attributed to `user_id`. + /// + /// Used by the user-memory wipe endpoint so that "wipe all memory" actually + /// wipes everything — the audit log is part of what the user is asking us + /// to forget. Returns the number of rows deleted. + /// + /// Uses the denormalized `session_extractions.user_id` column, so rows + /// orphaned by a prior session delete are still caught. The + /// JOIN-through-sessions clause is kept as a belt-and-suspenders + /// fallback in case a row was somehow inserted without the + /// denormalized column populated. + pub fn delete_user_extractions(&self, user_id: UserId) -> OpenFangResult { + let conn = self + .conn + .lock() + .map_err(|e| OpenFangError::Internal(e.to_string()))?; + let uid = user_id.0.to_string(); + let n = conn + .execute( + "DELETE FROM session_extractions \ + WHERE user_id = ?1 \ + OR session_id IN (SELECT id FROM sessions WHERE user_id = ?1)", + rusqlite::params![uid], + ) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + Ok(n) + } + + /// List extraction events for a user, newest first, with a `session_deleted` + /// flag derived from a LEFT JOIN against `sessions`. + /// + /// Returned tuple shape: `(extraction_id, session_id, agent_id, + /// created_at_rfc3339, session_deleted)`. Caller is the control API, + /// which serialises this to JSON. + pub fn list_user_extraction_audit( + &self, + user_id: UserId, + limit: usize, + ) -> OpenFangResult> { + let conn = self + .conn + .lock() + .map_err(|e| OpenFangError::Internal(e.to_string()))?; + let mut stmt = conn + .prepare( + "SELECT se.id, se.session_id, se.agent_id, se.created_at, \ + CASE WHEN s.id IS NULL THEN 1 ELSE 0 END AS deleted \ + FROM session_extractions se \ + LEFT JOIN sessions s ON s.id = se.session_id \ + WHERE se.user_id = ?1 \ + ORDER BY se.created_at DESC \ + LIMIT ?2", + ) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + let rows = stmt + .query_map( + rusqlite::params![user_id.0.to_string(), limit as i64], + |row| { + let id: String = row.get(0)?; + let sid: String = row.get(1)?; + let aid: String = row.get(2)?; + let created: String = row.get(3)?; + let deleted: i64 = row.get(4)?; + Ok((id, sid, aid, created, deleted != 0)) + }, + ) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + let mut out = Vec::new(); + for r in rows { + out.push(r.map_err(|e| OpenFangError::Memory(e.to_string()))?); + } + Ok(out) + } + + // ----------------------------------------------------------------- + // Structured memory: per-user topics + // ----------------------------------------------------------------- + + /// Upsert a user memory topic. + pub fn upsert_user_topic(&self, topic: &MemoryTopic) -> OpenFangResult<()> { + self.user_memory.upsert_topic(topic) + } + + /// Delete a single user memory topic (used by conflict resolution). + pub fn delete_user_topic(&self, user_id: UserId, topic: &str) -> OpenFangResult<()> { + self.user_memory.delete_topic(user_id, topic) + } + + /// Get the per-user topic index (no content) for session-start injection. + pub fn user_topic_index(&self, user_id: UserId) -> OpenFangResult> { + self.user_memory.get_index(user_id) + } + + /// Fetch full content for one user memory topic. `None` when missing or + /// expired. + pub fn user_topic(&self, user_id: UserId, topic: &str) -> OpenFangResult> { + self.user_memory.get_topic(user_id, topic) + } + + /// Store an embedding for a user topic. No-op when the topic doesn't exist. + pub fn store_user_topic_embedding( + &self, + user_id: UserId, + topic: &str, + embedding: &[f32], + ) -> OpenFangResult<()> { + self.user_memory.store_embedding(user_id, topic, embedding) + } + + /// Search user topics by cosine similarity against a query embedding. + pub fn search_user_topics_by_embedding( + &self, + user_id: UserId, + query: &[f32], + top_k: usize, + ) -> OpenFangResult> { + self.user_memory.search_by_embedding(user_id, query, top_k) + } + + /// Prune expired topics for a user (called opportunistically). + pub fn prune_expired_user_topics(&self, user_id: UserId) -> OpenFangResult { + self.user_memory.prune_expired(user_id) + } + + // ----------------------------------------------------------------- + // Structured memory: per-(user, agent) topics + // ----------------------------------------------------------------- + + /// Upsert a per-(user, agent) memory topic. + pub fn upsert_user_agent_topic(&self, topic: &UserAgentMemoryTopic) -> OpenFangResult<()> { + self.user_agent_memory.upsert_topic(topic) + } + + /// Get the per-(user, agent) topic index for session-start injection. + pub fn user_agent_topic_index( + &self, + user_id: UserId, + agent_id: AgentId, + ) -> OpenFangResult> { + self.user_agent_memory.get_index(user_id, agent_id) + } + + /// Fetch full content for one per-(user, agent) memory topic. + pub fn user_agent_topic( + &self, + user_id: UserId, + agent_id: AgentId, + topic: &str, + ) -> OpenFangResult> { + self.user_agent_memory.get_topic(user_id, agent_id, topic) + } + + /// Delete per-(user, agent) memory for one agent. Returns rows deleted. + pub fn delete_user_agent_memory( + &self, + user_id: UserId, + agent_id: AgentId, + ) -> OpenFangResult { + self.user_agent_memory + .delete_user_agent_memory(user_id, agent_id) + } + + // ----------------------------------------------------------------- + // Atomic user wipe (general + per-agent + audit) + // ----------------------------------------------------------------- + + /// Atomically wipe everything attributed to `user_id` across the three + /// structured-memory buckets. + /// + /// The three DELETEs run inside a single SQLite transaction so a partial + /// failure rolls back — callers never observe a half-wiped user. + /// The previous shape (three independent single-table deletes with + /// early-return on each error) could leave the user partially wiped if + /// the second or third call failed. Returns a per-bucket count so the + /// control-API handler can forward it verbatim. + pub fn wipe_user(&self, user_id: UserId) -> OpenFangResult { + let mut conn = self + .conn + .lock() + .map_err(|e| OpenFangError::Internal(e.to_string()))?; + let uid = user_id.0.to_string(); + let tx = conn + .transaction() + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + + let topics_deleted = tx + .execute( + "DELETE FROM user_memory_topics WHERE user_id = ?1", + rusqlite::params![uid], + ) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + let agent_topics_deleted = tx + .execute( + "DELETE FROM user_agent_memory_topics WHERE user_id = ?1", + rusqlite::params![uid], + ) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + let extractions_deleted = tx + .execute( + "DELETE FROM session_extractions \ + WHERE user_id = ?1 \ + OR session_id IN (SELECT id FROM sessions WHERE user_id = ?1)", + rusqlite::params![uid], + ) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + + tx.commit() + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + + Ok(WipeUserCounts { + topics_deleted, + agent_topics_deleted, + extractions_deleted, + }) + } + + /// Delete the session subtree rooted at `session_id` (the session itself + /// plus any descendants linked via `parent_session_id`). Returns the + /// number of sessions removed. + /// + /// Forks land in a later PR; this is wired here so the control API can + /// surface a single recursive delete now without coupling to forks. + pub fn delete_session_tree(&self, session_id: SessionId) -> OpenFangResult { + self.sessions.delete_session_tree(session_id) + } + // ----------------------------------------------------------------- // Paired devices persistence // ----------------------------------------------------------------- @@ -966,4 +1258,134 @@ mod tests { let claimed = substrate.task_claim("nobody").await.unwrap(); assert!(claimed.is_none()); } + + // ----------------------------------------------------------------- + // wipe_user — atomic per-user delete across the three structured-memory + // buckets. Must be scoped to the target user and never touch siblings. + // ----------------------------------------------------------------- + + fn seed_memory(substrate: &MemorySubstrate, user_id: UserId, agent_id: AgentId) { + substrate + .upsert_user_topic(&MemoryTopic { + user_id, + topic: "prefs".into(), + summary: "summary".into(), + content: "content".into(), + updated_at: chrono::Utc::now(), + expires_at: None, + }) + .unwrap(); + substrate + .upsert_user_agent_topic(&UserAgentMemoryTopic { + user_id, + agent_id, + topic: "with-jeeves".into(), + summary: "summary".into(), + content: "content".into(), + updated_at: chrono::Utc::now(), + }) + .unwrap(); + let session = substrate + .sessions + .create_session(agent_id, user_id) + .unwrap(); + substrate + .append_extraction( + session.id, + &SessionExtraction { + facts: vec!["a fact".into()], + ..Default::default() + }, + ) + .unwrap(); + } + + #[test] + fn test_wipe_user_returns_per_bucket_counts() { + let substrate = MemorySubstrate::open_in_memory(0.1).unwrap(); + let user = UserId::new(); + let agent = AgentId::new(); + seed_memory(&substrate, user, agent); + + let counts = substrate.wipe_user(user).unwrap(); + assert_eq!(counts.topics_deleted, 1); + assert_eq!(counts.agent_topics_deleted, 1); + assert_eq!(counts.extractions_deleted, 1); + + // Confirm the buckets are actually empty. + assert!(substrate.user_topic_index(user).unwrap().is_empty()); + assert!(substrate + .user_agent_topic_index(user, agent) + .unwrap() + .is_empty()); + } + + #[test] + fn test_wipe_user_is_scoped_to_target() { + // Wiping user A must not touch user B's data in any of the three + // buckets. + let substrate = MemorySubstrate::open_in_memory(0.1).unwrap(); + let agent = AgentId::new(); + let user_a = UserId::new(); + let user_b = UserId::new(); + + seed_memory(&substrate, user_a, agent); + seed_memory(&substrate, user_b, agent); + + let counts = substrate.wipe_user(user_a).unwrap(); + assert_eq!(counts.topics_deleted, 1); + assert_eq!(counts.agent_topics_deleted, 1); + assert_eq!(counts.extractions_deleted, 1); + + // User B still has everything. + assert_eq!(substrate.user_topic_index(user_b).unwrap().len(), 1); + assert_eq!( + substrate + .user_agent_topic_index(user_b, agent) + .unwrap() + .len(), + 1 + ); + let audit = substrate.list_user_extraction_audit(user_b, 10).unwrap(); + assert_eq!(audit.len(), 1); + } + + #[test] + fn test_wipe_user_idempotent_zero_counts() { + // Re-running wipe on an already-clean user returns 0/0/0 without + // erroring. + let substrate = MemorySubstrate::open_in_memory(0.1).unwrap(); + let user = UserId::new(); + let counts = substrate.wipe_user(user).unwrap(); + assert_eq!(counts.topics_deleted, 0); + assert_eq!(counts.agent_topics_deleted, 0); + assert_eq!(counts.extractions_deleted, 0); + } + + #[test] + fn test_audit_surfaces_session_deleted_flag() { + // Deleting the originating session must NOT lose the audit row — + // attribution survives via the denormalized user_id column, and the + // LEFT JOIN flips `session_deleted` to true. + let substrate = MemorySubstrate::open_in_memory(0.1).unwrap(); + let user = UserId::new(); + let agent = AgentId::new(); + let session = substrate.sessions.create_session(agent, user).unwrap(); + substrate + .append_extraction(session.id, &SessionExtraction::default()) + .unwrap(); + + // Before delete: session_deleted = false. + let audit = substrate.list_user_extraction_audit(user, 10).unwrap(); + assert_eq!(audit.len(), 1); + assert!(!audit[0].4, "session_deleted should start false"); + + // Delete the session. + substrate.sessions.delete_session(session.id).unwrap(); + + // After delete: extraction row still attributed to `user`, flag flips. + let audit = substrate.list_user_extraction_audit(user, 10).unwrap(); + assert_eq!(audit.len(), 1, "extraction must survive session delete"); + assert!(audit[0].4, "session_deleted should flip to true"); + } } diff --git a/crates/openfang-memory/src/user_agent_memory.rs b/crates/openfang-memory/src/user_agent_memory.rs new file mode 100644 index 0000000000..7e24dcd3c9 --- /dev/null +++ b/crates/openfang-memory/src/user_agent_memory.rs @@ -0,0 +1,386 @@ +//! Per-user-per-agent persistent memory organized into named topics. +//! +//! Stores memory specific to how a particular user interacts with a particular agent. +//! Separate from `user_memory` (general facts about a user across all agents). +//! +//! Examples: +//! - User Philippe + agent Jeeves: "Philippe wants Jeeves to use military time" +//! - User Philippe + agent google-mail-hand: "Philippe marks emails from acme.com as urgent" +//! +//! Dream writes these at session end. Injection reads the index at session start +//! and fetches relevant topic content per-message. + +use chrono::{DateTime, Utc}; +use openfang_types::agent::{AgentId, UserId}; +use openfang_types::error::{OpenFangError, OpenFangResult}; +use rusqlite::Connection; +use std::sync::{Arc, Mutex}; + +/// A single memory topic scoped to a (user, agent) pair. +#[derive(Debug, Clone)] +pub struct UserAgentMemoryTopic { + pub user_id: UserId, + pub agent_id: AgentId, + /// Topic key, e.g. "interaction_preferences", "delegated_tasks", "known_context". + pub topic: String, + /// One-line description shown in the session-start index. + pub summary: String, + /// Full topic content injected on per-turn retrieval. + pub content: String, + pub updated_at: DateTime, +} + +/// Lightweight index entry (no full content). +#[derive(Debug, Clone)] +pub struct UserAgentTopicIndexEntry { + pub topic: String, + pub summary: String, + pub updated_at: DateTime, +} + +/// Store for per-user-per-agent memory topics backed by SQLite. +#[derive(Clone)] +pub struct UserAgentMemoryStore { + conn: Arc>, +} + +impl UserAgentMemoryStore { + pub fn new(conn: Arc>) -> Self { + Self { conn } + } + + /// Upsert a topic — replaces existing content for this (user_id, agent_id, topic). + pub fn upsert_topic(&self, topic: &UserAgentMemoryTopic) -> OpenFangResult<()> { + let conn = self + .conn + .lock() + .map_err(|e| OpenFangError::Internal(e.to_string()))?; + conn.execute( + "INSERT INTO user_agent_memory_topics \ + (user_id, agent_id, topic, summary, content, updated_at) \ + VALUES (?1, ?2, ?3, ?4, ?5, ?6) \ + ON CONFLICT(user_id, agent_id, topic) \ + DO UPDATE SET summary = ?4, content = ?5, updated_at = ?6", + rusqlite::params![ + topic.user_id.0.to_string(), + topic.agent_id.0.to_string(), + topic.topic, + topic.summary, + topic.content, + topic.updated_at.to_rfc3339(), + ], + ) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + Ok(()) + } + + /// Get the index (topic names + summaries) for a (user, agent) pair — no full content. + pub fn get_index( + &self, + user_id: UserId, + agent_id: AgentId, + ) -> OpenFangResult> { + let conn = self + .conn + .lock() + .map_err(|e| OpenFangError::Internal(e.to_string()))?; + let mut stmt = conn + .prepare( + "SELECT topic, summary, updated_at FROM user_agent_memory_topics \ + WHERE user_id = ?1 AND agent_id = ?2 ORDER BY topic ASC", + ) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + + let rows = stmt + .query_map( + rusqlite::params![user_id.0.to_string(), agent_id.0.to_string()], + |row| { + let topic: String = row.get(0)?; + let summary: String = row.get(1)?; + let updated_at_str: String = row.get(2)?; + Ok((topic, summary, updated_at_str)) + }, + ) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + + let mut entries = Vec::new(); + for row in rows { + let (topic, summary, updated_at_str) = + row.map_err(|e| OpenFangError::Memory(e.to_string()))?; + let updated_at = updated_at_str + .parse::>() + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + entries.push(UserAgentTopicIndexEntry { + topic, + summary, + updated_at, + }); + } + Ok(entries) + } + + /// Get full content for a specific topic. + pub fn get_topic( + &self, + user_id: UserId, + agent_id: AgentId, + topic: &str, + ) -> OpenFangResult> { + let conn = self + .conn + .lock() + .map_err(|e| OpenFangError::Internal(e.to_string()))?; + let mut stmt = conn + .prepare( + "SELECT summary, content, updated_at FROM user_agent_memory_topics \ + WHERE user_id = ?1 AND agent_id = ?2 AND topic = ?3", + ) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + + let result = stmt.query_row( + rusqlite::params![user_id.0.to_string(), agent_id.0.to_string(), topic], + |row| { + let summary: String = row.get(0)?; + let content: String = row.get(1)?; + let updated_at_str: String = row.get(2)?; + Ok((summary, content, updated_at_str)) + }, + ); + + match result { + Ok((summary, content, updated_at_str)) => { + let updated_at = updated_at_str + .parse::>() + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + Ok(Some(UserAgentMemoryTopic { + user_id, + agent_id, + topic: topic.to_string(), + summary, + content, + updated_at, + })) + } + Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None), + Err(e) => Err(OpenFangError::Memory(e.to_string())), + } + } + + /// Delete all memory for a (user, agent) pair. + /// + /// Returns the number of rows deleted so the control API can report a + /// per-bucket count back to the caller, matching the shape used by the + /// other delete methods on this store. + pub fn delete_user_agent_memory( + &self, + user_id: UserId, + agent_id: AgentId, + ) -> OpenFangResult { + let conn = self + .conn + .lock() + .map_err(|e| OpenFangError::Internal(e.to_string()))?; + let n = conn + .execute( + "DELETE FROM user_agent_memory_topics WHERE user_id = ?1 AND agent_id = ?2", + rusqlite::params![user_id.0.to_string(), agent_id.0.to_string()], + ) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + Ok(n) + } + + /// Delete all memory across all agents for a user (e.g. on account deletion). + /// + /// Returns the number of rows deleted so the wipe endpoint can report a + /// per-bucket count back to the caller. + pub fn delete_all_user_memory(&self, user_id: UserId) -> OpenFangResult { + let conn = self + .conn + .lock() + .map_err(|e| OpenFangError::Internal(e.to_string()))?; + let n = conn + .execute( + "DELETE FROM user_agent_memory_topics WHERE user_id = ?1", + rusqlite::params![user_id.0.to_string()], + ) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + Ok(n) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::migration::run_migrations; + + fn setup() -> UserAgentMemoryStore { + let conn = Connection::open_in_memory().unwrap(); + run_migrations(&conn).unwrap(); + UserAgentMemoryStore::new(Arc::new(Mutex::new(conn))) + } + + fn make_topic( + user_id: UserId, + agent_id: AgentId, + topic: &str, + summary: &str, + content: &str, + ) -> UserAgentMemoryTopic { + UserAgentMemoryTopic { + user_id, + agent_id, + topic: topic.to_string(), + summary: summary.to_string(), + content: content.to_string(), + updated_at: Utc::now(), + } + } + + #[test] + fn test_upsert_and_get_topic() { + let store = setup(); + let user_id = UserId::new(); + let agent_id = AgentId::new(); + let t = make_topic( + user_id, + agent_id, + "interaction_preferences", + "How user prefers this agent to behave", + "Always use military time. Never suggest meetings before 9am.", + ); + store.upsert_topic(&t).unwrap(); + + let loaded = store + .get_topic(user_id, agent_id, "interaction_preferences") + .unwrap() + .unwrap(); + assert_eq!(loaded.topic, "interaction_preferences"); + assert_eq!(loaded.agent_id, agent_id); + assert_eq!(loaded.user_id, user_id); + } + + #[test] + fn test_upsert_replaces_existing() { + let store = setup(); + let user_id = UserId::new(); + let agent_id = AgentId::new(); + let t1 = make_topic(user_id, agent_id, "prefs", "v1", "Use formal tone."); + store.upsert_topic(&t1).unwrap(); + let t2 = make_topic(user_id, agent_id, "prefs", "v2", "Use casual tone."); + store.upsert_topic(&t2).unwrap(); + + let loaded = store + .get_topic(user_id, agent_id, "prefs") + .unwrap() + .unwrap(); + assert_eq!(loaded.summary, "v2"); + assert_eq!(loaded.content, "Use casual tone."); + } + + #[test] + fn test_get_index_returns_no_content() { + let store = setup(); + let user_id = UserId::new(); + let agent_id = AgentId::new(); + let t1 = make_topic(user_id, agent_id, "aaa", "Summary A", "Full content A"); + let t2 = make_topic(user_id, agent_id, "bbb", "Summary B", "Full content B"); + store.upsert_topic(&t1).unwrap(); + store.upsert_topic(&t2).unwrap(); + + let index = store.get_index(user_id, agent_id).unwrap(); + assert_eq!(index.len(), 2); + assert_eq!(index[0].topic, "aaa"); + assert_eq!(index[1].topic, "bbb"); + } + + #[test] + fn test_get_missing_topic_returns_none() { + let store = setup(); + let result = store + .get_topic(UserId::new(), AgentId::new(), "nonexistent") + .unwrap(); + assert!(result.is_none()); + } + + #[test] + fn test_different_users_isolated() { + let store = setup(); + let user_a = UserId::new(); + let user_b = UserId::new(); + let agent_id = AgentId::new(); + let t = make_topic(user_a, agent_id, "secret", "A's topic", "A's content"); + store.upsert_topic(&t).unwrap(); + + // user_b should not see user_a's topics for the same agent + let result = store.get_topic(user_b, agent_id, "secret").unwrap(); + assert!(result.is_none()); + let index = store.get_index(user_b, agent_id).unwrap(); + assert!(index.is_empty()); + } + + #[test] + fn test_different_agents_isolated() { + let store = setup(); + let user_id = UserId::new(); + let agent_a = AgentId::new(); + let agent_b = AgentId::new(); + let t = make_topic(user_id, agent_a, "prefs", "A's prefs", "Use formal tone."); + store.upsert_topic(&t).unwrap(); + + // Same user, different agent — should not see agent_a's topics + let result = store.get_topic(user_id, agent_b, "prefs").unwrap(); + assert!(result.is_none()); + let index = store.get_index(user_id, agent_b).unwrap(); + assert!(index.is_empty()); + } + + #[test] + fn test_same_topic_key_different_agents() { + let store = setup(); + let user_id = UserId::new(); + let agent_a = AgentId::new(); + let agent_b = AgentId::new(); + let ta = make_topic(user_id, agent_a, "prefs", "Agent A prefs", "Formal tone."); + let tb = make_topic(user_id, agent_b, "prefs", "Agent B prefs", "Casual tone."); + store.upsert_topic(&ta).unwrap(); + store.upsert_topic(&tb).unwrap(); + + let loaded_a = store.get_topic(user_id, agent_a, "prefs").unwrap().unwrap(); + let loaded_b = store.get_topic(user_id, agent_b, "prefs").unwrap().unwrap(); + assert_eq!(loaded_a.content, "Formal tone."); + assert_eq!(loaded_b.content, "Casual tone."); + } + + #[test] + fn test_delete_user_agent_memory() { + let store = setup(); + let user_id = UserId::new(); + let agent_id = AgentId::new(); + let t1 = make_topic(user_id, agent_id, "t1", "s1", "c1"); + let t2 = make_topic(user_id, agent_id, "t2", "s2", "c2"); + store.upsert_topic(&t1).unwrap(); + store.upsert_topic(&t2).unwrap(); + + store.delete_user_agent_memory(user_id, agent_id).unwrap(); + let index = store.get_index(user_id, agent_id).unwrap(); + assert!(index.is_empty()); + } + + #[test] + fn test_delete_all_user_memory_across_agents() { + let store = setup(); + let user_id = UserId::new(); + let agent_a = AgentId::new(); + let agent_b = AgentId::new(); + store + .upsert_topic(&make_topic(user_id, agent_a, "t1", "s", "c")) + .unwrap(); + store + .upsert_topic(&make_topic(user_id, agent_b, "t1", "s", "c")) + .unwrap(); + + store.delete_all_user_memory(user_id).unwrap(); + assert!(store.get_index(user_id, agent_a).unwrap().is_empty()); + assert!(store.get_index(user_id, agent_b).unwrap().is_empty()); + } +} diff --git a/crates/openfang-memory/src/user_memory.rs b/crates/openfang-memory/src/user_memory.rs new file mode 100644 index 0000000000..f1936f5413 --- /dev/null +++ b/crates/openfang-memory/src/user_memory.rs @@ -0,0 +1,579 @@ +//! Per-user persistent memory organized into named topics. +//! +//! Dream writes topic entries at session end; injection reads the index +//! (topic names + summaries) at session start and fetches full topic content +//! per-message for relevant topics. +//! +//! Topics can optionally carry an `expires_at` timestamp for time-sensitive facts +//! (travel plans, temporary preferences). Expired topics are pruned at read time. +//! The `embedding` column stores a packed little-endian f32 vector for optional +//! cosine-similarity retrieval; it is populated asynchronously after write and +//! may be NULL for topics created before embedding was configured. + +use chrono::{DateTime, Utc}; +use openfang_types::agent::UserId; +use openfang_types::error::{OpenFangError, OpenFangResult}; +use rusqlite::Connection; +use std::sync::{Arc, Mutex}; + +/// A single memory topic for a user. +#[derive(Debug, Clone)] +pub struct MemoryTopic { + pub user_id: UserId, + /// Topic key, e.g. "work_context", "preferences", "open_items". + pub topic: String, + /// One-line description shown in the session-start index. + pub summary: String, + /// Full topic content injected on per-turn retrieval. + pub content: String, + pub updated_at: DateTime, + /// Optional expiry for time-sensitive facts. Expired topics are invisible at read time. + pub expires_at: Option>, +} + +/// Lightweight index entry for a topic (no full content). +#[derive(Debug, Clone)] +pub struct TopicIndexEntry { + pub topic: String, + pub summary: String, + pub updated_at: DateTime, +} + +/// User memory store backed by SQLite. +#[derive(Clone)] +pub struct UserMemoryStore { + conn: Arc>, +} + +impl UserMemoryStore { + pub fn new(conn: Arc>) -> Self { + Self { conn } + } + + /// Upsert a topic — replaces existing content for this (user_id, topic) pair. + pub fn upsert_topic(&self, topic: &MemoryTopic) -> OpenFangResult<()> { + let conn = self + .conn + .lock() + .map_err(|e| OpenFangError::Internal(e.to_string()))?; + conn.execute( + "INSERT INTO user_memory_topics (user_id, topic, summary, content, updated_at, expires_at) \ + VALUES (?1, ?2, ?3, ?4, ?5, ?6) \ + ON CONFLICT(user_id, topic) DO UPDATE SET \ + summary = ?3, content = ?4, updated_at = ?5, expires_at = ?6", + rusqlite::params![ + topic.user_id.0.to_string(), + topic.topic, + topic.summary, + topic.content, + topic.updated_at.to_rfc3339(), + topic.expires_at.as_ref().map(|d| d.to_rfc3339()), + ], + ) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + Ok(()) + } + + /// Delete a single topic for a user. Used by conflict resolution (supersedes). + pub fn delete_topic(&self, user_id: UserId, topic: &str) -> OpenFangResult<()> { + let conn = self + .conn + .lock() + .map_err(|e| OpenFangError::Internal(e.to_string()))?; + conn.execute( + "DELETE FROM user_memory_topics WHERE user_id = ?1 AND topic = ?2", + rusqlite::params![user_id.0.to_string(), topic], + ) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + Ok(()) + } + + /// Get the index (topic names + summaries) for a user — no full content. + /// Expired topics are excluded. + pub fn get_index(&self, user_id: UserId) -> OpenFangResult> { + let conn = self + .conn + .lock() + .map_err(|e| OpenFangError::Internal(e.to_string()))?; + let mut stmt = conn + .prepare( + "SELECT topic, summary, updated_at FROM user_memory_topics \ + WHERE user_id = ?1 \ + AND (expires_at IS NULL OR datetime(expires_at) > datetime('now')) \ + ORDER BY topic ASC", + ) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + + let rows = stmt + .query_map(rusqlite::params![user_id.0.to_string()], |row| { + let topic: String = row.get(0)?; + let summary: String = row.get(1)?; + let updated_at_str: String = row.get(2)?; + Ok((topic, summary, updated_at_str)) + }) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + + let mut entries = Vec::new(); + for row in rows { + let (topic, summary, updated_at_str) = + row.map_err(|e| OpenFangError::Memory(e.to_string()))?; + let updated_at = updated_at_str + .parse::>() + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + entries.push(TopicIndexEntry { + topic, + summary, + updated_at, + }); + } + Ok(entries) + } + + /// Get full content for a specific topic. Returns None if the topic does not exist + /// or has expired. + pub fn get_topic(&self, user_id: UserId, topic: &str) -> OpenFangResult> { + let conn = self + .conn + .lock() + .map_err(|e| OpenFangError::Internal(e.to_string()))?; + let mut stmt = conn + .prepare( + "SELECT summary, content, updated_at, expires_at FROM user_memory_topics \ + WHERE user_id = ?1 AND topic = ?2 \ + AND (expires_at IS NULL OR datetime(expires_at) > datetime('now'))", + ) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + + let result = stmt.query_row(rusqlite::params![user_id.0.to_string(), topic], |row| { + let summary: String = row.get(0)?; + let content: String = row.get(1)?; + let updated_at_str: String = row.get(2)?; + let expires_at_str: Option = row.get(3)?; + Ok((summary, content, updated_at_str, expires_at_str)) + }); + + match result { + Ok((summary, content, updated_at_str, expires_at_str)) => { + let updated_at = updated_at_str + .parse::>() + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + let expires_at = expires_at_str + .map(|s| { + s.parse::>() + .map_err(|e| OpenFangError::Memory(e.to_string())) + }) + .transpose()?; + Ok(Some(MemoryTopic { + user_id, + topic: topic.to_string(), + summary, + content, + updated_at, + expires_at, + })) + } + Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None), + Err(e) => Err(OpenFangError::Memory(e.to_string())), + } + } + + /// Store a pre-computed embedding for a topic. + /// + /// Embeddings are packed as little-endian f32 bytes. Called asynchronously + /// after `upsert_topic` when an embedding provider is configured. + pub fn store_embedding( + &self, + user_id: UserId, + topic: &str, + embedding: &[f32], + ) -> OpenFangResult<()> { + let blob = pack_embedding(embedding); + let conn = self + .conn + .lock() + .map_err(|e| OpenFangError::Internal(e.to_string()))?; + conn.execute( + "UPDATE user_memory_topics SET embedding = ?1 WHERE user_id = ?2 AND topic = ?3", + rusqlite::params![blob, user_id.0.to_string(), topic], + ) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + Ok(()) + } + + /// Search topics by cosine similarity to a query embedding. + /// + /// Returns topic names sorted by descending similarity, limited to `top_k` results. + /// Topics without an embedding are skipped. Falls back to an empty list if no + /// embeddings are stored (caller should use the LLM selector instead). + pub fn search_by_embedding( + &self, + user_id: UserId, + query: &[f32], + top_k: usize, + ) -> OpenFangResult> { + let conn = self + .conn + .lock() + .map_err(|e| OpenFangError::Internal(e.to_string()))?; + let mut stmt = conn + .prepare( + "SELECT topic, embedding FROM user_memory_topics \ + WHERE user_id = ?1 AND embedding IS NOT NULL \ + AND (expires_at IS NULL OR datetime(expires_at) > datetime('now'))", + ) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + + let rows = stmt + .query_map(rusqlite::params![user_id.0.to_string()], |row| { + let topic: String = row.get(0)?; + let blob: Vec = row.get(1)?; + Ok((topic, blob)) + }) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + + let mut scored: Vec<(String, f32)> = Vec::new(); + for row in rows { + let (topic, blob) = row.map_err(|e| OpenFangError::Memory(e.to_string()))?; + let emb = unpack_embedding(&blob); + let score = cosine_similarity(query, &emb); + scored.push((topic, score)); + } + + scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); + Ok(scored.into_iter().take(top_k).map(|(t, _)| t).collect()) + } + + /// Delete all memory for a user. + /// + /// Returns the number of rows deleted so that callers (e.g. the wipe + /// endpoint) can report a per-bucket count back to the user. + pub fn delete_user_memory(&self, user_id: UserId) -> OpenFangResult { + let conn = self + .conn + .lock() + .map_err(|e| OpenFangError::Internal(e.to_string()))?; + let n = conn + .execute( + "DELETE FROM user_memory_topics WHERE user_id = ?1", + rusqlite::params![user_id.0.to_string()], + ) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + Ok(n) + } + + /// Delete expired topics for a user. Called opportunistically — expiry is also + /// enforced at read time, but periodic pruning keeps the table tidy. + pub fn prune_expired(&self, user_id: UserId) -> OpenFangResult { + let conn = self + .conn + .lock() + .map_err(|e| OpenFangError::Internal(e.to_string()))?; + let n = conn + .execute( + "DELETE FROM user_memory_topics \ + WHERE user_id = ?1 AND expires_at IS NOT NULL AND datetime(expires_at) <= datetime('now')", + rusqlite::params![user_id.0.to_string()], + ) + .map_err(|e| OpenFangError::Memory(e.to_string()))?; + Ok(n) + } +} + +// ── Embedding helpers ──────────────────────────────────────────────────────── + +/// Pack f32 slice to little-endian bytes. +fn pack_embedding(v: &[f32]) -> Vec { + let mut out = Vec::with_capacity(v.len() * 4); + for &x in v { + out.extend_from_slice(&x.to_le_bytes()); + } + out +} + +/// Unpack little-endian bytes to f32 slice. +fn unpack_embedding(bytes: &[u8]) -> Vec { + bytes + .chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect() +} + +/// Cosine similarity between two equal-length vectors. Returns 0.0 on zero vectors. +pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 { + if a.len() != b.len() || a.is_empty() { + return 0.0; + } + let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum(); + let na: f32 = a.iter().map(|x| x * x).sum::().sqrt(); + let nb: f32 = b.iter().map(|x| x * x).sum::().sqrt(); + if na == 0.0 || nb == 0.0 { + 0.0 + } else { + dot / (na * nb) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::migration::run_migrations; + + fn setup() -> UserMemoryStore { + let conn = Connection::open_in_memory().unwrap(); + run_migrations(&conn).unwrap(); + UserMemoryStore::new(Arc::new(Mutex::new(conn))) + } + + fn make_topic(user_id: UserId, topic: &str, summary: &str, content: &str) -> MemoryTopic { + MemoryTopic { + user_id, + topic: topic.to_string(), + summary: summary.to_string(), + content: content.to_string(), + updated_at: Utc::now(), + expires_at: None, + } + } + + #[test] + fn test_upsert_and_get_topic() { + let store = setup(); + let user_id = UserId::new(); + let t = make_topic(user_id, "work_context", "Work summary", "I work at Acme."); + store.upsert_topic(&t).unwrap(); + + let loaded = store.get_topic(user_id, "work_context").unwrap().unwrap(); + assert_eq!(loaded.topic, "work_context"); + assert_eq!(loaded.summary, "Work summary"); + assert_eq!(loaded.content, "I work at Acme."); + assert!(loaded.expires_at.is_none()); + } + + #[test] + fn test_upsert_replaces_existing() { + let store = setup(); + let user_id = UserId::new(); + let t1 = make_topic(user_id, "preferences", "Prefs v1", "I prefer tea."); + store.upsert_topic(&t1).unwrap(); + + let t2 = make_topic(user_id, "preferences", "Prefs v2", "I prefer coffee."); + store.upsert_topic(&t2).unwrap(); + + let loaded = store.get_topic(user_id, "preferences").unwrap().unwrap(); + assert_eq!(loaded.summary, "Prefs v2"); + assert_eq!(loaded.content, "I prefer coffee."); + } + + #[test] + fn test_get_index_returns_no_content() { + let store = setup(); + let user_id = UserId::new(); + let t1 = make_topic(user_id, "aaa", "Summary A", "Full content A"); + let t2 = make_topic(user_id, "bbb", "Summary B", "Full content B"); + store.upsert_topic(&t1).unwrap(); + store.upsert_topic(&t2).unwrap(); + + let index = store.get_index(user_id).unwrap(); + assert_eq!(index.len(), 2); + assert_eq!(index[0].topic, "aaa"); + assert_eq!(index[0].summary, "Summary A"); + assert_eq!(index[1].topic, "bbb"); + assert_eq!(index[1].summary, "Summary B"); + } + + #[test] + fn test_get_missing_topic_returns_none() { + let store = setup(); + let user_id = UserId::new(); + let result = store.get_topic(user_id, "nonexistent").unwrap(); + assert!(result.is_none()); + } + + #[test] + fn test_different_users_isolated() { + let store = setup(); + let user_a = UserId::new(); + let user_b = UserId::new(); + let t = make_topic(user_a, "private", "A's summary", "A's content"); + store.upsert_topic(&t).unwrap(); + + let result = store.get_topic(user_b, "private").unwrap(); + assert!(result.is_none()); + + let index = store.get_index(user_b).unwrap(); + assert!(index.is_empty()); + } + + #[test] + fn test_delete_user_memory() { + let store = setup(); + let user_id = UserId::new(); + let t1 = make_topic(user_id, "t1", "s1", "c1"); + let t2 = make_topic(user_id, "t2", "s2", "c2"); + store.upsert_topic(&t1).unwrap(); + store.upsert_topic(&t2).unwrap(); + + store.delete_user_memory(user_id).unwrap(); + + let index = store.get_index(user_id).unwrap(); + assert!(index.is_empty()); + } + + #[test] + fn test_delete_single_topic() { + let store = setup(); + let user_id = UserId::new(); + store + .upsert_topic(&make_topic(user_id, "keep", "keep", "keep")) + .unwrap(); + store + .upsert_topic(&make_topic(user_id, "remove", "remove", "remove")) + .unwrap(); + + store.delete_topic(user_id, "remove").unwrap(); + + let index = store.get_index(user_id).unwrap(); + assert_eq!(index.len(), 1); + assert_eq!(index[0].topic, "keep"); + } + + #[test] + fn test_expired_topic_invisible() { + let store = setup(); + let user_id = UserId::new(); + // Topic that expired 1 hour ago + let past = Utc::now() - chrono::Duration::hours(1); + let t = MemoryTopic { + user_id, + topic: "expired".to_string(), + summary: "old".to_string(), + content: "stale content".to_string(), + updated_at: Utc::now(), + expires_at: Some(past), + }; + store.upsert_topic(&t).unwrap(); + + // Should not be returned by get_topic or get_index + assert!(store.get_topic(user_id, "expired").unwrap().is_none()); + assert!(store.get_index(user_id).unwrap().is_empty()); + } + + #[test] + fn test_non_expired_topic_visible() { + let store = setup(); + let user_id = UserId::new(); + let future = Utc::now() + chrono::Duration::days(7); + let t = MemoryTopic { + user_id, + topic: "travel".to_string(), + summary: "upcoming trip".to_string(), + content: "Paris trip on 2026-04-10".to_string(), + updated_at: Utc::now(), + expires_at: Some(future), + }; + store.upsert_topic(&t).unwrap(); + + let loaded = store.get_topic(user_id, "travel").unwrap(); + assert!(loaded.is_some()); + let index = store.get_index(user_id).unwrap(); + assert_eq!(index.len(), 1); + } + + #[test] + fn test_prune_expired() { + let store = setup(); + let user_id = UserId::new(); + let past = Utc::now() - chrono::Duration::hours(1); + let future = Utc::now() + chrono::Duration::days(1); + + store + .upsert_topic(&MemoryTopic { + user_id, + topic: "expired".into(), + summary: "".into(), + content: "".into(), + updated_at: Utc::now(), + expires_at: Some(past), + }) + .unwrap(); + store + .upsert_topic(&MemoryTopic { + user_id, + topic: "valid".into(), + summary: "".into(), + content: "".into(), + updated_at: Utc::now(), + expires_at: Some(future), + }) + .unwrap(); + store + .upsert_topic(&make_topic(user_id, "permanent", "", "")) + .unwrap(); + + let pruned = store.prune_expired(user_id).unwrap(); + assert_eq!(pruned, 1); + + // Permanent and future-expiry remain + let conn = store.conn.lock().unwrap(); + let count: i64 = conn + .query_row( + "SELECT COUNT(*) FROM user_memory_topics WHERE user_id = ?1", + rusqlite::params![user_id.0.to_string()], + |r| r.get(0), + ) + .unwrap(); + assert_eq!(count, 2); + } + + #[test] + fn test_cosine_similarity() { + let a = vec![1.0f32, 0.0, 0.0]; + let b = vec![1.0f32, 0.0, 0.0]; + assert!((cosine_similarity(&a, &b) - 1.0).abs() < 1e-5); + + let c = vec![0.0f32, 1.0, 0.0]; + assert!((cosine_similarity(&a, &c)).abs() < 1e-5); + } + + #[test] + fn test_embedding_store_and_search() { + let store = setup(); + let user_id = UserId::new(); + store + .upsert_topic(&make_topic( + user_id, + "rust", + "Rust programming", + "loves Rust", + )) + .unwrap(); + store + .upsert_topic(&make_topic( + user_id, + "python", + "Python programming", + "also uses Python", + )) + .unwrap(); + + // Embed "rust" topic with [1,0,0] and "python" with [0,1,0] + store + .store_embedding(user_id, "rust", &[1.0, 0.0, 0.0]) + .unwrap(); + store + .store_embedding(user_id, "python", &[0.0, 1.0, 0.0]) + .unwrap(); + + // Query closest to [1,0,0] — should be "rust" + let results = store + .search_by_embedding(user_id, &[1.0, 0.0, 0.0], 1) + .unwrap(); + assert_eq!(results, vec!["rust".to_string()]); + + // Query closest to [0,1,0] — should be "python" + let results = store + .search_by_embedding(user_id, &[0.0, 1.0, 0.0], 1) + .unwrap(); + assert_eq!(results, vec!["python".to_string()]); + } +} diff --git a/crates/openfang-runtime/src/prompt_builder.rs b/crates/openfang-runtime/src/prompt_builder.rs index dd28ae396e..41a8bfb00d 100644 --- a/crates/openfang-runtime/src/prompt_builder.rs +++ b/crates/openfang-runtime/src/prompt_builder.rs @@ -66,6 +66,15 @@ pub struct PromptContext { /// Read per-turn by the kernel so external writers (cron jobs, integrations) /// are reflected in the next LLM call. See issue #843. pub context_md: Option, + /// User memory index — topic names + summaries for the session owner. + /// + /// `Some(formatted_block)` when the agent has opted in to structured + /// memory (`MemoryConfig::is_structured()`) AND the user has at least + /// one non-empty topic; `None` for default-memory agents and for empty + /// indexes. Injected near the top of the system prompt so the model + /// can decide whether to fetch full content for any topic during the + /// turn. + pub user_memory_context: Option, } /// Build the complete system prompt from a `PromptContext`. @@ -152,6 +161,23 @@ pub fn build_system_prompt(ctx: &PromptContext) -> String { sections.push(build_user_section(ctx.user_name.as_deref())); } + // Section 8.5 — User Memory Index (skip for subagents, only when present) + // + // Populated by the kernel only for agents that opted in to structured + // memory via `[memory] system = "structured"`. Default agents never + // produce a value here, so the section is skipped entirely. + if !ctx.is_subagent { + if let Some(ref mem_ctx) = ctx.user_memory_context { + if !mem_ctx.trim().is_empty() { + sections.push(format!( + "## What I Remember About You\n\ + The following topics are stored in your personal memory. \ + You can recall any topic for more detail during the conversation.\n\n{mem_ctx}" + )); + } + } + } + // Section 9 — Channel Awareness (skip for subagents) if !ctx.is_subagent { if let Some(ref channel) = ctx.channel_type { diff --git a/crates/openfang-types/src/agent.rs b/crates/openfang-types/src/agent.rs index 5589b1e35e..dbac6acb62 100644 --- a/crates/openfang-types/src/agent.rs +++ b/crates/openfang-types/src/agent.rs @@ -516,6 +516,61 @@ pub struct AgentManifest { /// `agent_send` results stay focused. See issue #871. #[serde(default)] pub max_history_messages: Option, + /// Per-agent memory system selection. Omitted in TOML → defaults to + /// `MemorySystem::Summarization` (current OpenFang behavior). Set to + /// `structured` to enable the LLM-driven structured memory pipeline + /// (extractions, dream consolidation, context injection). The default + /// is skipped on serialization so manifests that don't opt in stay + /// clean. + #[serde(default, skip_serializing_if = "MemoryConfig::is_default")] + pub memory: MemoryConfig, +} + +/// Per-agent memory system selection. +/// +/// Agents opt in to advanced memory features here. The default +/// (`Summarization`) preserves existing OpenFang behavior — LLM summarization +/// of older messages when the message count or token threshold is exceeded, +/// with no structured fact extraction and no background dream consolidation. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum MemorySystem { + /// Current OpenFang behavior: LLM summarization of old messages when + /// message count or token threshold is exceeded. No structured extraction, + /// no dreamer. + #[default] + Summarization, + /// Structured memory: LLM extracts facts/preferences/decisions/tasks/ + /// open_items during compaction and on overflow drain, with periodic + /// dream consolidation during inactivity. Opt-in per agent. The producer + /// (extraction + dreamer) lands in a later PR — this PR only wires the + /// storage layer and the gate. + Structured, +} + +/// Per-agent memory configuration. +/// +/// All fields default so existing agent manifests that omit `[memory]` +/// deserialize unchanged — preserving upstream behavior. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(default)] +pub struct MemoryConfig { + /// Memory system to use for this agent. + pub system: MemorySystem, +} + +impl MemoryConfig { + /// Returns true if the structured memory system is enabled for this agent. + pub fn is_structured(&self) -> bool { + matches!(self.system, MemorySystem::Structured) + } + + /// Returns true when this config equals the default. Used by + /// `skip_serializing_if` on `AgentManifest::memory` to keep round-trip + /// TOML clean for the no-opt-in case. + pub fn is_default(&self) -> bool { + *self == Self::default() + } } /// Runtime default for `AgentManifest::max_history_messages` when the agent @@ -569,6 +624,7 @@ impl Default for AgentManifest { tool_blocklist: Vec::new(), cache_context: false, max_history_messages: None, + memory: MemoryConfig::default(), } } } @@ -829,6 +885,7 @@ mod tests { tool_blocklist: Vec::new(), cache_context: false, max_history_messages: None, + memory: MemoryConfig::default(), }; let json = serde_json::to_string(&manifest).unwrap(); let deserialized: AgentManifest = serde_json::from_str(&json).unwrap(); @@ -1347,4 +1404,100 @@ memory_write = ["self.*"] vec!["self.*".to_string()] ); } + + // --- Per-agent memory system selection --- + + #[test] + fn test_manifest_memory_defaults_to_summarization() { + // No [memory] block → default to Summarization (current upstream behavior). + let toml_str = r#" +name = "no-memory-block" + +[model] +provider = "groq" +model = "llama-3.3-70b-versatile" +system_prompt = "hi" +"#; + let manifest: AgentManifest = toml::from_str(toml_str).unwrap(); + assert_eq!(manifest.memory.system, MemorySystem::Summarization); + assert!(!manifest.memory.is_structured()); + } + + #[test] + fn test_manifest_memory_opt_in_to_structured() { + // Explicit [memory] system = "structured" turns on the extraction+dream path. + let toml_str = r#" +name = "structured-memory-agent" + +[model] +provider = "groq" +model = "llama-3.3-70b-versatile" +system_prompt = "hi" + +[memory] +system = "structured" +"#; + let manifest: AgentManifest = toml::from_str(toml_str).unwrap(); + assert_eq!(manifest.memory.system, MemorySystem::Structured); + assert!(manifest.memory.is_structured()); + } + + #[test] + fn test_manifest_memory_explicit_summarization() { + // Explicit summarization should also parse and equal the default. + let toml_str = r#" +name = "summ-agent" + +[model] +provider = "groq" +model = "llama-3.3-70b-versatile" +system_prompt = "hi" + +[memory] +system = "summarization" +"#; + let manifest: AgentManifest = toml::from_str(toml_str).unwrap(); + assert_eq!(manifest.memory.system, MemorySystem::Summarization); + } + + #[test] + fn test_memory_config_default_is_summarization() { + let cfg = MemoryConfig::default(); + assert_eq!(cfg.system, MemorySystem::Summarization); + assert!(!cfg.is_structured()); + assert!(cfg.is_default()); + } + + #[test] + fn test_memory_config_skip_serializing_default() { + // Manifest with default memory config should not include [memory] in + // TOML output, so existing manifests round-trip clean. + let mut manifest = AgentManifest { + name: "round-trip".into(), + model: ModelConfig { + provider: "groq".into(), + model: "llama-3.3-70b-versatile".into(), + system_prompt: "hi".into(), + ..Default::default() + }, + ..Default::default() + }; + let out = toml::to_string(&manifest).unwrap(); + assert!( + !out.contains("[memory]"), + "default MemoryConfig should be skipped in TOML output, got:\n{out}" + ); + + // Opt in → [memory] reappears. + manifest.memory.system = MemorySystem::Structured; + let out = toml::to_string(&manifest).unwrap(); + assert!( + out.contains("[memory]"), + "structured opt-in should serialize [memory]; got:\n{out}" + ); + assert!( + out.contains("system = \"structured\""), + "structured opt-in should emit `system = \"structured\"`; got:\n{out}" + ); + } } From 1726b0abc355be83bd947e54300f20278ec176de Mon Sep 17 00:00:00 2001 From: Philippe Branchu Date: Wed, 3 Jun 2026 06:43:04 +0000 Subject: [PATCH 04/12] =?UTF-8?q?feat(memory):=20structured=20memory=20pro?= =?UTF-8?q?ducer=20=E2=80=94=20extract=20+=20mini=5Fdream=20+=20dreamer?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the per-user memory producer code that populates the storage layer introduced in PR 2. Entirely gated behind `manifest.memory.is_structured()` — default (Summarization) agents see zero behaviour change, while opted-in agents pay a structured-extraction LLM call when the context overflows and a dream-consolidation pass when their session goes idle. Components: * `compactor::extract_structured()` — LLM-driven extraction producing `SessionExtraction { facts, preferences, decisions, tasks, open_items }`. Filters out `MessageSource::ContextInjection` so calendar/email summaries do not bleed into long-term memory. Falls back to the existing extraction (or an empty one) on LLM error or persistent parse failure — never propagates an error to the agent turn. * `compactor::needs_extraction()` + `count_tool_calls()` — trigger gates on tokens-since-last or tool-calls-since-last thresholds. * `compactor::SessionExtraction` re-export — runtime callers reach the struct without pulling `openfang-memory` directly. * `context_overflow::overflow_drain_count()` — peek at how many leading messages overflow recovery would drain, without applying the trim, so mini-dream can extract from them first. * `mini_dream` module — in-loop consumer that, immediately before the overflow recovery trims messages, runs extract + dream and persists the resulting topics into the user memory store. Non-fatal: errors are logged, never returned. * `dreamer` module — session-end consolidation pass: merges all `SessionExtraction` records accumulated during compaction plus the recent message tail into 3–7 topic-organised user memory entries, with conflict resolution (`supersedes` list) and expiry tagging. * `agent_loop.rs` integration — two `if manifest.memory.is_structured()` call sites (streaming + non-streaming) wire mini_dream into the iteration prologue. Default agents skip entirely. * `OpenFangKernel::trigger_session_dream()` + `run_session_lifecycle_loop()` — background loop polls `agent_last_active` every 30 s; for each agent idle longer than `[sessions] gap_secs`, fires the dream pass. Per-agent gate: structured-memory only. Per-session gate: skip if no real user activity (pure `[AUTONOMOUS TICK]` / `[SCHEDULED TICK]` / `ContextInjection` sessions are no-ops — the production thundering-herd fix from commit `aa4ec5c` on `branchu`). * `SessionsConfig` (in `KernelConfig`) — `[sessions] gap_secs` knob, default 300 s. Set to 0 to disable the dreamer loop entirely. * `Message::context_injection()` helper — small constructor for the new `ContextInjection`-tagged messages used by extract/dream tests and the activity-gating predicate. Tests (12 new, all green): * `compactor::test_extract_structured_filters_context_injections` * `compactor::test_extract_structured_fallback_on_empty_response` * `compactor::test_needs_extraction_token_threshold` * `compactor::test_needs_extraction_tool_calls_threshold` * `compactor::test_count_tool_calls` * `dreamer::test_dream_filters_context_injections` * `dreamer::test_dream_result_has_topics` * `dreamer::test_dream_conflict_resolution` * `dreamer::test_dream_expiry_tagging` * `dreamer::test_dream_fallback_on_parse_error` * `mini_dream::test_structured_memory_gate_skips_default_agents` * `mini_dream::test_structured_memory_gate_fires_for_opted_in_agents` Built on top of `pr/memory-storage` (PR 2). No UI, route, schema, or `MemorySystem`/`MemoryConfig` changes — that surface area is PR 2's. Co-Authored-By: Claude Sonnet 4.6 --- crates/openfang-kernel/src/kernel.rs | 306 +++++++++++ crates/openfang-runtime/src/agent_loop.rs | 51 +- crates/openfang-runtime/src/compactor.rs | 315 ++++++++++- .../openfang-runtime/src/context_overflow.rs | 32 ++ crates/openfang-runtime/src/dreamer.rs | 499 ++++++++++++++++++ crates/openfang-runtime/src/lib.rs | 2 + crates/openfang-runtime/src/mini_dream.rs | 192 +++++++ crates/openfang-types/src/config.rs | 34 ++ crates/openfang-types/src/message.rs | 16 + 9 files changed, 1444 insertions(+), 3 deletions(-) create mode 100644 crates/openfang-runtime/src/dreamer.rs create mode 100644 crates/openfang-runtime/src/mini_dream.rs diff --git a/crates/openfang-kernel/src/kernel.rs b/crates/openfang-kernel/src/kernel.rs index d21104149d..29ae01a8c4 100644 --- a/crates/openfang-kernel/src/kernel.rs +++ b/crates/openfang-kernel/src/kernel.rs @@ -178,6 +178,10 @@ pub struct OpenFangKernel { /// session corruption when multiple messages arrive concurrently (e.g. rapid voice /// messages via Telegram). Different agents can still run in parallel. agent_msg_locks: dashmap::DashMap>>, + /// Last-activity timestamp per agent — populated on every successful agent + /// loop completion. The session lifecycle loop polls this to detect idle + /// sessions and trigger dream consolidation for structured-memory agents. + agent_last_active: dashmap::DashMap, /// Weak self-reference for trigger dispatch (set after Arc wrapping). self_handle: OnceLock>, } @@ -1304,6 +1308,7 @@ impl OpenFangKernel { default_model_override: std::sync::RwLock::new(None), fallback_providers_override: std::sync::RwLock::new(None), agent_msg_locks: dashmap::DashMap::new(), + agent_last_active: dashmap::DashMap::new(), self_handle: OnceLock::new(), }; @@ -2027,6 +2032,10 @@ impl OpenFangKernel { // Record token usage for quota tracking self.scheduler.record_usage(agent_id, &result.total_usage); + // Update last-active time for the session lifecycle / dream trigger. + self.agent_last_active + .insert(agent_id, std::time::Instant::now()); + // Update last active time let _ = self.registry.set_state(agent_id, AgentState::Running); @@ -2127,6 +2136,9 @@ impl OpenFangKernel { kernel_clone .scheduler .record_usage(agent_id, &result.total_usage); + kernel_clone + .agent_last_active + .insert(agent_id, std::time::Instant::now()); let _ = kernel_clone .registry .set_state(agent_id, AgentState::Running); @@ -2538,6 +2550,10 @@ impl OpenFangKernel { tool_calls: result.iterations.saturating_sub(1), }); + kernel_clone + .agent_last_active + .insert(agent_id, std::time::Instant::now()); + let _ = kernel_clone .registry .set_state(agent_id, AgentState::Running); @@ -4590,6 +4606,23 @@ impl OpenFangKernel { // Start heartbeat monitor for agent health checking self.start_heartbeat_monitor(); + // Session lifecycle: inactivity detection and dream consolidation. + // + // The dreamer background loop runs kernel-wide; per-agent gating happens + // inside `trigger_session_dream` via `manifest.memory.is_structured()`, + // so default (Summarization) agents are no-ops even when their last-active + // timestamp expires. Set `[sessions] gap_secs = 0` to disable entirely. + if self.config.sessions.gap_secs > 0 { + let kernel = Arc::clone(self); + tokio::spawn(async move { + kernel.run_session_lifecycle_loop().await; + }); + info!( + gap_secs = self.config.sessions.gap_secs, + "Session lifecycle monitor started" + ); + } + // Start OFP peer node if network is enabled if self.config.network_enabled && !self.config.network.shared_secret.is_empty() { let kernel = Arc::clone(self); @@ -5112,6 +5145,279 @@ impl OpenFangKernel { info!("Heartbeat monitor started (interval: {}s)", interval_secs); } + /// Background loop: check for inactive agent sessions and trigger dream consolidation. + /// + /// Wakes up every 30 seconds. For each agent that has been idle longer than + /// `sessions.gap_secs`, triggers the dream pass (structured extraction → + /// user memory consolidation) and removes the agent from the active map. + /// + /// The per-agent `manifest.memory.is_structured()` gate inside + /// `trigger_session_dream` ensures default agents are no-ops even if they + /// briefly show up in `agent_last_active` from a turn that ran before the + /// opt-in check. + async fn run_session_lifecycle_loop(self: &Arc) { + let check_interval = std::time::Duration::from_secs(30); + let inactivity_threshold = std::time::Duration::from_secs(self.config.sessions.gap_secs); + + let mut interval = tokio::time::interval(check_interval); + interval.tick().await; // Skip initial tick + + loop { + interval.tick().await; + if self.supervisor.is_shutting_down() { + break; + } + + let now = std::time::Instant::now(); + let expired: Vec = self + .agent_last_active + .iter() + .filter(|entry| now.duration_since(*entry.value()) >= inactivity_threshold) + .map(|entry| *entry.key()) + .collect(); + + for agent_id in expired { + self.agent_last_active.remove(&agent_id); + let kernel = Arc::clone(self); + tokio::spawn(async move { + kernel.trigger_session_dream(agent_id).await; + }); + } + } + } + + /// Run the dream pass for an agent session that has gone idle. + /// + /// Loads the session's structured extractions (accumulated during compaction passes), + /// plus the recent message tail, and consolidates into topic-organized user memory. + /// + /// Gates: + /// - The agent must opt in to structured memory (`manifest.memory.is_structured()`). + /// - The session must contain *real* user activity — pure autonomous-tick sessions + /// (where the only inputs are `[AUTONOMOUS TICK]` / `[SCHEDULED TICK]` prompts or + /// `ContextInjection` messages) are skipped to avoid the thundering-herd problem + /// that surfaced in production when every heartbeat fired a dream. + pub async fn trigger_session_dream(self: &Arc, agent_id: AgentId) { + use openfang_memory::user_memory::MemoryTopic; + use openfang_runtime::compactor::{extract_structured, CompactionConfig}; + use openfang_runtime::dreamer::dream; + + let entry = match self.registry.get(agent_id) { + Some(e) => e, + None => return, + }; + + // Gate: dream consolidation only runs for agents that have opted in + // to structured memory. Agents on the default Summarization path never + // accumulate extractions, so there is nothing to consolidate. + if !entry.manifest.memory.is_structured() { + debug!( + agent_id = %agent_id, + "Dream: skipping — agent memory.system is not structured" + ); + return; + } + + let session = match self.memory.get_session(entry.session_id) { + Ok(Some(s)) => s, + _ => return, + }; + + if session.messages.is_empty() { + return; + } + + // Skip dream if the session had no real user activity — only autonomous + // tick prompts or context injections. A pure-tick session (heartbeat fires, + // agent responds NO_REPLY, nothing else) is not worth consolidating. + // This is the activity-gating fix from production (commit aa4ec5c on branchu). + let has_real_activity = session.messages.iter().any(|msg| { + if msg.role != openfang_types::message::Role::User { + return false; + } + if msg.source == Some(openfang_types::message::MessageSource::ContextInjection) { + return false; + } + let text = msg.content.text_content(); + !text.starts_with("[AUTONOMOUS TICK]") && !text.starts_with("[SCHEDULED TICK]") + }); + + if !has_real_activity { + debug!( + agent_id = %agent_id, + messages = session.messages.len(), + "Dream: skipping — no real user activity in session" + ); + return; + } + + let user_id = session.user_id; + let session_id = session.id; + + info!( + agent_id = %agent_id, + session_id = %session_id, + messages = session.messages.len(), + "Dream: session idle, starting consolidation" + ); + + // Run a final extraction pass if needed + let config = CompactionConfig::default(); + let existing_extractions = self.memory.load_extractions(session_id).unwrap_or_default(); + + // Always run a final extraction to capture anything since last compaction + let final_extraction = if !session.messages.is_empty() { + let existing = existing_extractions.last().cloned(); + match extract_structured( + self.resolve_driver(&entry.manifest) + .ok() + .unwrap_or_else(|| { + Arc::new(crate::kernel::StubDriver) + as Arc + }), + &entry.manifest.model.model, + &session.messages, + existing.as_ref(), + &config, + ) + .await + { + Ok(extraction) => { + let _ = self.memory.append_extraction(session_id, &extraction); + extraction + } + Err(e) => { + warn!(agent_id = %agent_id, "Dream: final extraction failed: {e}"); + existing_extractions.last().cloned().unwrap_or_default() + } + } + } else { + existing_extractions.last().cloned().unwrap_or_default() + }; + + // Reload all extractions including the new one + let all_extractions = self.memory.load_extractions(session_id).unwrap_or_default(); + + if all_extractions.is_empty() + && final_extraction.facts.is_empty() + && final_extraction.preferences.is_empty() + { + debug!(agent_id = %agent_id, "Dream: no content to consolidate, skipping"); + return; + } + + // Load existing user memory topics (name + full content) for conflict detection. + // The dream prompt uses these to detect contradictions and emit a `supersedes` list. + let existing_topics: Vec<(String, String)> = self + .memory + .user_topic_index(user_id) + .unwrap_or_default() + .into_iter() + .filter_map(|entry| { + self.memory + .user_topic(user_id, &entry.topic) + .ok() + .flatten() + .map(|t| (entry.topic, t.content)) + }) + .collect(); + + let driver = match self.resolve_driver(&entry.manifest) { + Ok(d) => d, + Err(e) => { + warn!(agent_id = %agent_id, "Dream: no driver available: {e}"); + return; + } + }; + + // Take the last ~20 messages as the "recent tail" (not compacted) + let recent_tail: Vec<_> = session + .messages + .iter() + .rev() + .take(20) + .cloned() + .rev() + .collect(); + + match dream( + driver, + &entry.manifest.model.model, + &all_extractions, + &recent_tail, + user_id, + &existing_topics, + &config, + ) + .await + { + Ok(result) => { + info!( + agent_id = %agent_id, + topics = result.topics.len(), + superseded = result.superseded_topics.len(), + "Dream: consolidation complete, persisting topics" + ); + + // Delete superseded topics before writing new ones (conflict resolution). + for superseded in &result.superseded_topics { + if let Err(e) = self.memory.delete_user_topic(user_id, superseded) { + warn!(agent_id = %agent_id, topic = %superseded, "Dream: failed to delete superseded topic: {e}"); + } else { + info!(agent_id = %agent_id, topic = %superseded, "Dream: deleted superseded topic"); + } + } + + // Write new topics, propagating expires_at from dream output. + let now = chrono::Utc::now(); + for dt in &result.topics { + let expires_at = dt + .expires_at + .as_deref() + .and_then(|s| s.parse::>().ok()); + let topic = MemoryTopic { + user_id, + topic: dt.topic.clone(), + summary: dt.summary.clone(), + content: dt.content.clone(), + updated_at: now, + expires_at, + }; + if let Err(e) = self.memory.upsert_user_topic(&topic) { + warn!(agent_id = %agent_id, topic = %dt.topic, "Dream: failed to persist topic: {e}"); + continue; + } + // Embed the topic content for semantic retrieval. + if let Some(emb) = &self.embedding_driver { + match emb.embed_one(&dt.content).await { + Ok(vec) => { + if let Err(e) = self + .memory + .store_user_topic_embedding(user_id, &dt.topic, &vec) + { + warn!(agent_id = %agent_id, topic = %dt.topic, "Dream: failed to store embedding: {e}"); + } + } + Err(e) => { + warn!(agent_id = %agent_id, topic = %dt.topic, "Dream: embedding failed: {e}"); + } + } + } + } + + // Prune any expired topics left over from previous sessions. + if let Ok(n) = self.memory.prune_expired_user_topics(user_id) { + if n > 0 { + info!(agent_id = %agent_id, pruned = n, "Dream: pruned expired topics"); + } + } + } + Err(e) => { + warn!(agent_id = %agent_id, "Dream: consolidation failed: {e}"); + } + } + } + /// Start the background loop / register triggers for a single agent. pub fn start_background_for_agent( self: &Arc, diff --git a/crates/openfang-runtime/src/agent_loop.rs b/crates/openfang-runtime/src/agent_loop.rs index 6f29244fc6..e6bc5df091 100644 --- a/crates/openfang-runtime/src/agent_loop.rs +++ b/crates/openfang-runtime/src/agent_loop.rs @@ -4,8 +4,9 @@ //! calling the LLM, executing tool calls, and saving the conversation. use crate::auth_cooldown::{CooldownVerdict, ProviderCooldown}; +use crate::compactor::CompactionConfig; use crate::context_budget::{apply_context_guard, truncate_tool_result_dynamic, ContextBudget}; -use crate::context_overflow::{recover_from_overflow, RecoveryStage}; +use crate::context_overflow::{overflow_drain_count, recover_from_overflow, RecoveryStage}; use crate::embedding::EmbeddingDriver; use crate::kernel_handle::KernelHandle; use crate::llm_driver::{CompletionRequest, DriverConfig, LlmDriver, LlmError, StreamEvent}; @@ -506,6 +507,30 @@ pub async fn run_agent_loop( for iteration in 0..max_iterations { debug!(iteration, "Agent loop iteration"); + // Mini-dream: before trimming, extract facts from messages about to be discarded + // so they are preserved in user memory rather than lost. + // + // Gated by the per-agent memory system selection: only runs when the + // agent has explicitly opted in to structured memory. Agents on the + // default (Summarization) path skip extraction entirely. + if manifest.memory.is_structured() { + let drain = + overflow_drain_count(&messages, &system_prompt, available_tools, ctx_window); + if drain > 0 { + let config = CompactionConfig::default(); + crate::mini_dream::run_mini_dream( + &messages[..drain], + session.user_id, + Arc::clone(&driver), + &manifest.model.model, + memory, + &config, + embedding_driver.map(|e| e as &(dyn EmbeddingDriver + Send + Sync)), + ) + .await; + } + } + // Context overflow recovery pipeline (replaces emergency_trim_messages) let recovery = recover_from_overflow(&mut messages, &system_prompt, available_tools, ctx_window); @@ -1725,6 +1750,30 @@ pub async fn run_agent_loop_streaming( for iteration in 0..max_iterations { debug!(iteration, "Streaming agent loop iteration"); + // Mini-dream: before trimming, extract facts from messages about to be discarded + // so they are preserved in user memory rather than lost. + // + // Gated by the per-agent memory system selection: only runs when the + // agent has explicitly opted in to structured memory. Agents on the + // default (Summarization) path skip extraction entirely. + if manifest.memory.is_structured() { + let drain = + overflow_drain_count(&messages, &system_prompt, available_tools, ctx_window); + if drain > 0 { + let config = CompactionConfig::default(); + crate::mini_dream::run_mini_dream( + &messages[..drain], + session.user_id, + Arc::clone(&driver), + &manifest.model.model, + memory, + &config, + embedding_driver.map(|e| e as &(dyn EmbeddingDriver + Send + Sync)), + ) + .await; + } + } + // Context overflow recovery pipeline (replaces emergency_trim_messages) let recovery = recover_from_overflow(&mut messages, &system_prompt, available_tools, ctx_window); diff --git a/crates/openfang-runtime/src/compactor.rs b/crates/openfang-runtime/src/compactor.rs index 2b50a236c1..669143d590 100644 --- a/crates/openfang-runtime/src/compactor.rs +++ b/crates/openfang-runtime/src/compactor.rs @@ -13,12 +13,18 @@ use crate::llm_driver::{CompletionRequest, LlmDriver}; use crate::str_utils::safe_truncate_str; use openfang_memory::session::Session; -use openfang_types::message::{ContentBlock, Message, MessageContent, Role}; +use openfang_types::message::{ContentBlock, Message, MessageContent, MessageSource, Role}; use openfang_types::tool::ToolDefinition; use serde::Serialize; use std::sync::Arc; use tracing::{info, warn}; +// Re-export so callers can use `openfang_runtime::compactor::SessionExtraction` +// without pulling in openfang-memory directly. The struct itself lives in +// openfang-memory (PR 2) because the storage layer needs it; runtime callers +// reach it through this re-export. +pub use openfang_memory::session::SessionExtraction; + /// Configuration for session compaction. #[derive(Debug, Clone)] pub struct CompactionConfig { @@ -44,6 +50,9 @@ pub struct CompactionConfig { pub token_threshold_ratio: f64, /// Model context window size in tokens. pub context_window_tokens: usize, + /// Number of tool calls between structured extraction passes. + /// Used by `needs_extraction` to gate the mini-dream / extractor flow. + pub tool_calls_between_extractions: usize, } impl Default for CompactionConfig { @@ -60,6 +69,7 @@ impl Default for CompactionConfig { max_retries: 3, token_threshold_ratio: 0.7, context_window_tokens: 200_000, + tool_calls_between_extractions: 10, } } } @@ -84,6 +94,160 @@ pub fn needs_compaction(session: &Session, config: &CompactionConfig) -> bool { session.messages.len() > config.threshold } +/// Check whether a session needs structured extraction. +/// Triggers on token count OR tool call count since last extraction. +pub fn needs_extraction( + _messages_since_last: usize, + tokens_since_last: usize, + tool_calls_since_last: usize, + config: &CompactionConfig, +) -> bool { + let token_threshold = (config.context_window_tokens as f64 * 0.3) as usize; + tokens_since_last > token_threshold + || tool_calls_since_last >= config.tool_calls_between_extractions +} + +/// Count the number of tool calls (ToolUse blocks) in a slice of messages. +pub fn count_tool_calls(messages: &[Message]) -> usize { + messages + .iter() + .map(|msg| match &msg.content { + MessageContent::Blocks(blocks) => blocks + .iter() + .filter(|b| matches!(b, ContentBlock::ToolUse { .. })) + .count(), + _ => 0, + }) + .sum() +} + +/// Extract structured memory from a slice of messages using the LLM. +/// +/// Produces facts, preferences, decisions, tasks, and open items. +/// Skips messages tagged as context injections (MessageSource::ContextInjection) +/// so calendar / email summaries do not get re-extracted into long-term memory. +/// +/// On LLM error or persistent parse failure, returns the `existing` extraction +/// (or `SessionExtraction::default()`) rather than propagating an error — the +/// extraction path is non-critical and must never fail the agent turn. +pub async fn extract_structured( + driver: Arc, + model: &str, + messages: &[Message], + existing: Option<&SessionExtraction>, + config: &CompactionConfig, +) -> Result { + // Filter out context injection messages + let filtered: Vec<&Message> = messages + .iter() + .filter(|m| m.source != Some(MessageSource::ContextInjection)) + .collect(); + + let owned: Vec = filtered.into_iter().cloned().collect(); + let conversation_text = build_conversation_text(&owned, config); + + let existing_section = if let Some(ext) = existing { + format!( + "Current memory state (update this, don't just append):\n\ + Facts: {}\n\ + Preferences: {}\n\ + Decisions: {}\n\ + Tasks: {}\n\ + Open items: {}\n\n", + ext.facts.join("; "), + ext.preferences.join("; "), + ext.decisions.join("; "), + ext.tasks.join("; "), + ext.open_items.join("; "), + ) + } else { + String::new() + }; + + let prompt = format!( + "You are analyzing a conversation to extract structured memory for long-term storage.\n\n\ + {existing_section}\ + New conversation to integrate:\n\ + ---\n\ + {conversation_text}\ + ---\n\n\ + Extract or update the following. Each item should be a concise, self-contained statement.\n\ + Return ONLY valid JSON with this exact structure:\n\ + {{\n\ + \"facts\": [\"...\"],\n\ + \"preferences\": [\"...\"],\n\ + \"decisions\": [\"...\"],\n\ + \"tasks\": [\"...\"],\n\ + \"open_items\": [\"...\"]\n\ + }}\n\n\ + Where:\n\ + - facts: Facts stated or confirmed about the user or their world\n\ + - preferences: User preferences, working style, communication preferences\n\ + - decisions: Decisions made or conclusions reached\n\ + - tasks: Tasks completed or delegated\n\ + - open_items: Unresolved questions or pending actions" + ); + + let request = CompletionRequest { + model: model.to_string(), + messages: vec![Message { + role: Role::User, + content: MessageContent::Blocks(vec![ContentBlock::Text { + text: prompt, + provider_metadata: None, + }]), + source: None, + ..Default::default() + }], + tools: vec![], + max_tokens: config.max_summary_tokens, + temperature: 0.3, + system: Some( + "You are a memory extraction assistant. Extract structured information from conversations \ + and return ONLY valid JSON. Do not include any text outside the JSON object." + .to_string(), + ), + thinking: None, + }; + + let fallback = || existing.cloned().unwrap_or_default(); + + for attempt in 0..config.max_retries { + match driver.complete(request.clone()).await { + Ok(response) => { + let text = response.text(); + if text.is_empty() { + warn!(attempt, "Empty response from LLM for structured extraction"); + continue; + } + // Strip markdown code fences if present + let json_str = text + .trim() + .trim_start_matches("```json") + .trim_start_matches("```") + .trim_end_matches("```") + .trim(); + match serde_json::from_str::(json_str) { + Ok(extraction) => { + info!("Structured extraction complete"); + return Ok(extraction); + } + Err(e) => { + warn!(attempt, error = %e, "Failed to parse structured extraction JSON, retrying"); + } + } + } + Err(e) => { + warn!(attempt, error = %e, "LLM call failed during structured extraction"); + } + } + } + + // All retries failed — return existing or empty rather than erroring + warn!("Structured extraction failed after all retries, returning fallback"); + Ok(fallback()) +} + /// Estimate token count for a set of messages, optional system prompt, and tool definitions. /// /// Uses the chars/4 heuristic — not exact, but good enough for budget gating. @@ -324,7 +488,7 @@ fn is_oversized(message: &Message, config: &CompactionConfig) -> bool { /// /// Handles all content block types: text, tool use, tool result, image, unknown. /// Oversized messages are truncated inline with a marker. -fn build_conversation_text(messages: &[Message], config: &CompactionConfig) -> String { +pub fn build_conversation_text(messages: &[Message], config: &CompactionConfig) -> String { let mut conversation_text = String::new(); for msg in messages { @@ -1529,4 +1693,151 @@ mod tests { assert_eq!(adjust_split_for_tool_pairs(&messages, 1), 1); assert_eq!(adjust_split_for_tool_pairs(&messages, 5), 5); } + + // --- Structured extraction tests (PR 3) --- + + #[test] + fn test_extract_structured_filters_context_injections() { + use openfang_types::message::MessageSource; + + // Build a mix of normal and context-injection messages + let messages = vec![ + Message::user("I prefer dark mode."), + Message::context_injection("Here is your calendar: meeting at 3pm"), + Message::assistant("Noted!"), + Message::context_injection("Email summary: 5 unread emails"), + ]; + + // Count tool_calls is 0 for these messages + assert_eq!(count_tool_calls(&messages), 0); + + // Context injections should be filtered out; verify by checking that + // the context_injection messages are tagged correctly + let filtered: Vec<&Message> = messages + .iter() + .filter(|m| m.source != Some(MessageSource::ContextInjection)) + .collect(); + assert_eq!( + filtered.len(), + 2, + "Should filter out 2 context injection messages" + ); + assert!(filtered[0].content.text_content().contains("dark mode")); + assert!(filtered[1].content.text_content().contains("Noted")); + } + + #[test] + fn test_needs_extraction_token_threshold() { + let config = CompactionConfig { + context_window_tokens: 100_000, + ..CompactionConfig::default() + }; + // 30% of 100_000 = 30_000 + assert!(!needs_extraction(0, 29_999, 0, &config)); + assert!(needs_extraction(0, 30_001, 0, &config)); + } + + #[test] + fn test_needs_extraction_tool_calls_threshold() { + let config = CompactionConfig { + tool_calls_between_extractions: 10, + ..CompactionConfig::default() + }; + assert!(!needs_extraction(0, 0, 9, &config)); + assert!(needs_extraction(0, 0, 10, &config)); + assert!(needs_extraction(0, 0, 15, &config)); + } + + #[test] + fn test_count_tool_calls() { + let messages = vec![ + Message::user("hello"), + Message { + msg_id: uuid::Uuid::new_v4().to_string(), + provider_msg_id: None, + role: Role::Assistant, + content: MessageContent::Blocks(vec![ + ContentBlock::Text { + text: "searching".to_string(), + provider_metadata: None, + }, + ContentBlock::ToolUse { + id: "t1".to_string(), + name: "web_search".to_string(), + input: serde_json::json!({}), + provider_metadata: None, + }, + ]), + source: None, + }, + Message { + msg_id: uuid::Uuid::new_v4().to_string(), + provider_msg_id: None, + role: Role::Assistant, + content: MessageContent::Blocks(vec![ContentBlock::ToolUse { + id: "t2".to_string(), + name: "read_file".to_string(), + input: serde_json::json!({}), + provider_metadata: None, + }]), + source: None, + }, + Message::assistant("Done."), + ]; + assert_eq!(count_tool_calls(&messages), 2); + } + + #[tokio::test] + async fn test_extract_structured_fallback_on_empty_response() { + use crate::llm_driver::{CompletionResponse, LlmError}; + use async_trait::async_trait; + + struct EmptyDriver; + + #[async_trait] + impl LlmDriver for EmptyDriver { + async fn complete( + &self, + _req: CompletionRequest, + ) -> Result { + Ok(CompletionResponse { + content: vec![ContentBlock::Text { + text: "this is not valid json {{{".to_string(), + provider_metadata: None, + }], + stop_reason: openfang_types::message::StopReason::EndTurn, + tool_calls: vec![], + usage: TokenUsage { + input_tokens: 10, + output_tokens: 5, + }, + }) + } + } + + let messages = vec![Message::user("I like Rust.")]; + let existing = SessionExtraction { + facts: vec!["User likes Rust".to_string()], + ..SessionExtraction::default() + }; + + let config = CompactionConfig { + max_retries: 2, + ..CompactionConfig::default() + }; + + // Should fall back to existing extraction without erroring + let result = extract_structured( + Arc::new(EmptyDriver), + "test-model", + &messages, + Some(&existing), + &config, + ) + .await + .unwrap(); + + // Returns the existing extraction as fallback + assert_eq!(result.facts, vec!["User likes Rust".to_string()]); + } } diff --git a/crates/openfang-runtime/src/context_overflow.rs b/crates/openfang-runtime/src/context_overflow.rs index 3d9fa8c48f..6dcefd16fc 100644 --- a/crates/openfang-runtime/src/context_overflow.rs +++ b/crates/openfang-runtime/src/context_overflow.rs @@ -111,6 +111,38 @@ fn estimate_tokens(messages: &[Message], system_prompt: &str, tools: &[ToolDefin crate::compactor::estimate_token_count(messages, Some(system_prompt), Some(tools)) } +/// Returns the number of leading messages that overflow recovery would drain, +/// without actually applying the trim. +/// +/// 0 means no recovery needed. The caller can run mini-dream on `messages[..n]` +/// before calling `recover_from_overflow` to preserve facts before discard. +pub fn overflow_drain_count( + messages: &[Message], + system_prompt: &str, + tools: &[ToolDefinition], + context_window: usize, +) -> usize { + let estimated = estimate_tokens(messages, system_prompt, tools); + let threshold_70 = (context_window as f64 * 0.70) as usize; + let threshold_90 = (context_window as f64 * 0.90) as usize; + + if estimated <= threshold_70 { + return 0; + } + + if estimated <= threshold_90 { + // Matches Stage 1: keep last 10 + let keep = 10.min(messages.len()); + let raw_remove = messages.len().saturating_sub(keep); + return safe_drain_boundary(messages, raw_remove); + } + + // Matches Stage 2: keep last 4 + let keep = 4.min(messages.len()); + let raw_remove = messages.len().saturating_sub(keep); + safe_drain_boundary(messages, raw_remove) +} + /// Run the 4-stage overflow recovery pipeline. /// /// Returns the recovery stage applied and the number of messages/results affected. diff --git a/crates/openfang-runtime/src/dreamer.rs b/crates/openfang-runtime/src/dreamer.rs new file mode 100644 index 0000000000..70c4bc21c6 --- /dev/null +++ b/crates/openfang-runtime/src/dreamer.rs @@ -0,0 +1,499 @@ +//! Dream: semantic consolidation of a conversation session into user memory. +//! +//! Runs at session end (triggered by inactivity timeout). Takes all structured +//! extractions accumulated during compaction passes, plus the recent message tail, +//! and consolidates them into topic-organized user memory. +//! +//! Runs as a background task — does not block the next message. +//! Context-injection messages (calendar/email summaries) are excluded. +//! +//! ## Conflict resolution +//! The dream prompt includes the full content of existing memory topics so the LLM +//! can detect contradictions. It returns a `supersedes` list of topic names that +//! should be deleted before the new topics are written. +//! +//! ## Expiry +//! Time-sensitive topics (travel plans, temporary preferences) are tagged with an +//! optional `expires_at` ISO timestamp. Durable facts carry `null`. + +use crate::compactor::{build_conversation_text, CompactionConfig, SessionExtraction}; +use crate::llm_driver::{CompletionRequest, LlmDriver}; +use openfang_types::agent::UserId; +use openfang_types::message::{ContentBlock, Message, MessageContent, MessageSource, Role}; +use std::sync::Arc; +use tracing::{info, warn}; + +/// A topic-organized memory entry produced by dream. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct DreamTopic { + pub topic: String, + pub summary: String, // one-liner for the index + pub content: String, // full content for retrieval + /// ISO 8601 timestamp after which this topic expires and is invisible at read time. + /// Null for durable facts. + #[serde(default)] + pub expires_at: Option, +} + +/// Result of a dream pass. +#[derive(Debug, Clone)] +pub struct DreamResult { + pub topics: Vec, + pub user_id: UserId, + /// Topic names from existing memory that are superseded by this dream pass + /// and should be deleted before the new topics are written. + pub superseded_topics: Vec, +} + +/// Run the dream pass: consolidate extractions + recent messages into topic memory. +/// +/// - `extractions`: all SessionExtraction records from compaction during this session +/// - `recent_messages`: recent tail messages kept verbatim (already filtered) +/// - `user_id`: whose memory this belongs to +/// - `existing_topics`: full existing user memory `(topic_name, content)` pairs for +/// conflict detection. Dream may supersede stale entries. +pub async fn dream( + driver: Arc, + model: &str, + extractions: &[SessionExtraction], + recent_messages: &[Message], + user_id: UserId, + existing_topics: &[(String, String)], + config: &CompactionConfig, +) -> Result { + // Filter recent_messages to exclude ContextInjection + let filtered_recent: Vec = recent_messages + .iter() + .filter(|m| m.source != Some(MessageSource::ContextInjection)) + .cloned() + .collect(); + + let recent_conversation = build_conversation_text(&filtered_recent, config); + + // Merge all extractions into consolidated lists + let all_facts: Vec = extractions.iter().flat_map(|e| e.facts.clone()).collect(); + let all_preferences: Vec = extractions + .iter() + .flat_map(|e| e.preferences.clone()) + .collect(); + let all_decisions: Vec = extractions + .iter() + .flat_map(|e| e.decisions.clone()) + .collect(); + let all_tasks: Vec = extractions.iter().flat_map(|e| e.tasks.clone()).collect(); + let all_open_items: Vec = extractions + .iter() + .flat_map(|e| e.open_items.clone()) + .collect(); + + let fmt_list = |items: &[String]| -> String { + if items.is_empty() { + "(none)".to_string() + } else { + items.join("; ") + } + }; + + // Build existing topics section with full content for conflict detection + let existing_topics_str = if existing_topics.is_empty() { + "(none)".to_string() + } else { + existing_topics + .iter() + .map(|(name, content)| format!(" [{name}]: {content}")) + .collect::>() + .join("\n") + }; + + let prompt = format!( + "You are consolidating a conversation into persistent user memory organized by topic.\n\n\ + Accumulated insights from this session:\n\ + Facts: {facts}\n\ + Preferences: {prefs}\n\ + Decisions: {decisions}\n\ + Tasks completed: {tasks}\n\ + Open items: {open}\n\n\ + Recent conversation:\n\ + {recent}\n\n\ + Existing memory topics (full content — check for contradictions):\n\ + {existing}\n\n\ + Instructions:\n\ + 1. Produce 3-7 named memory topics that reflect the complete, current state of the user's memory.\n\ + 2. If new information contradicts an existing topic, the new information wins. List the old \ + topic name in `supersedes`.\n\ + 3. For time-sensitive facts (travel plans, temporary preferences, deadlines), set `expires_at` \ + to an ISO 8601 timestamp (e.g. \"2026-04-10T00:00:00Z\"). For durable facts set `expires_at` to null.\n\ + 4. Topic names must be snake_case. Good examples: work_context, preferences, open_items, \ + family_context, upcoming_travel.\n\n\ + Return ONLY valid JSON:\n\ + {{\n\ + \"supersedes\": [\"old_topic_name_if_any\"],\n\ + \"topics\": [\n\ + {{\n\ + \"topic\": \"snake_case_name\",\n\ + \"summary\": \"one line description of what this topic contains\",\n\ + \"content\": \"full content for this topic, 2-10 sentences\",\n\ + \"expires_at\": null\n\ + }}\n\ + ]\n\ + }}", + facts = fmt_list(&all_facts), + prefs = fmt_list(&all_preferences), + decisions = fmt_list(&all_decisions), + tasks = fmt_list(&all_tasks), + open = fmt_list(&all_open_items), + recent = recent_conversation, + existing = existing_topics_str, + ); + + let request = CompletionRequest { + model: model.to_string(), + messages: vec![Message { + role: Role::User, + content: MessageContent::Blocks(vec![ContentBlock::Text { + text: prompt, + provider_metadata: None, + }]), + source: None, + ..Default::default() + }], + tools: vec![], + max_tokens: config.max_summary_tokens * 2, // dreamer can produce more content + temperature: 0.3, + system: Some( + "You are a memory consolidation assistant. Organize conversation insights into \ + structured topics with conflict resolution and expiry tagging. \ + Return ONLY valid JSON. Do not include any text outside the JSON object." + .to_string(), + ), + thinking: None, + }; + + for attempt in 0..config.max_retries { + match driver.complete(request.clone()).await { + Ok(response) => { + let text = response.text(); + if text.is_empty() { + warn!(attempt, "Empty response from LLM for dream pass"); + continue; + } + // Strip markdown code fences if present + let json_str = text + .trim() + .trim_start_matches("```json") + .trim_start_matches("```") + .trim_end_matches("```") + .trim(); + match parse_dream_response(json_str) { + Ok((topics, superseded)) => { + info!( + topics = topics.len(), + superseded = superseded.len(), + "Dream pass complete" + ); + return Ok(DreamResult { + topics, + user_id, + superseded_topics: superseded, + }); + } + Err(e) => { + warn!(attempt, error = %e, "Failed to parse dream response JSON, retrying"); + } + } + } + Err(e) => { + warn!(attempt, error = %e, "LLM call failed during dream pass"); + } + } + } + + // Fallback: produce a single "general" topic from concatenated extractions + warn!("Dream pass failed after all retries, falling back to general topic"); + let general_content = build_fallback_content( + &all_facts, + &all_preferences, + &all_decisions, + &all_tasks, + &all_open_items, + ); + let fallback_topic = DreamTopic { + topic: "general".to_string(), + summary: "General session memory (dream consolidation unavailable)".to_string(), + content: general_content, + expires_at: None, + }; + Ok(DreamResult { + topics: vec![fallback_topic], + user_id, + superseded_topics: vec![], + }) +} + +/// Parse the LLM response JSON into (topics, superseded_topics). +fn parse_dream_response(json_str: &str) -> Result<(Vec, Vec), String> { + #[derive(serde::Deserialize)] + struct DreamResponse { + #[serde(default)] + supersedes: Vec, + topics: Vec, + } + serde_json::from_str::(json_str) + .map(|r| (r.topics, r.supersedes)) + .map_err(|e| e.to_string()) +} + +/// Build a fallback "general" topic content from all extraction lists. +fn build_fallback_content( + facts: &[String], + preferences: &[String], + decisions: &[String], + tasks: &[String], + open_items: &[String], +) -> String { + let mut parts = Vec::new(); + if !facts.is_empty() { + parts.push(format!("Facts: {}", facts.join("; "))); + } + if !preferences.is_empty() { + parts.push(format!("Preferences: {}", preferences.join("; "))); + } + if !decisions.is_empty() { + parts.push(format!("Decisions: {}", decisions.join("; "))); + } + if !tasks.is_empty() { + parts.push(format!("Tasks: {}", tasks.join("; "))); + } + if !open_items.is_empty() { + parts.push(format!("Open items: {}", open_items.join("; "))); + } + if parts.is_empty() { + "No structured information extracted from this session.".to_string() + } else { + parts.join("\n") + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::compactor::CompactionConfig; + use crate::llm_driver::{CompletionResponse, LlmError}; + use async_trait::async_trait; + use openfang_types::message::{StopReason, TokenUsage}; + use std::sync::Arc; + + struct MockDreamDriver { + response: String, + } + + #[async_trait] + impl LlmDriver for MockDreamDriver { + async fn complete(&self, _req: CompletionRequest) -> Result { + Ok(CompletionResponse { + content: vec![ContentBlock::Text { + text: self.response.clone(), + provider_metadata: None, + }], + stop_reason: StopReason::EndTurn, + tool_calls: vec![], + usage: TokenUsage { + input_tokens: 100, + output_tokens: 50, + }, + }) + } + } + + #[test] + fn test_dream_filters_context_injections() { + let messages = [ + Message::user("What is 2+2?"), + Message::context_injection("Calendar: meeting at 3pm"), + Message::assistant("It's 4."), + Message::context_injection("Email: 3 unread"), + ]; + + let filtered: Vec = messages + .iter() + .filter(|m| m.source != Some(MessageSource::ContextInjection)) + .cloned() + .collect(); + + assert_eq!(filtered.len(), 2); + assert!(filtered[0].content.text_content().contains("2+2")); + assert!(filtered[1].content.text_content().contains("It's 4")); + } + + #[tokio::test] + async fn test_dream_result_has_topics() { + let valid_response = r#"{ + "supersedes": [], + "topics": [ + { + "topic": "preferences", + "summary": "User interface and work preferences", + "content": "The user prefers dark mode. They like concise responses.", + "expires_at": null + }, + { + "topic": "open_items", + "summary": "Pending tasks and questions", + "content": "There is a pending code review. The user wants to set up CI/CD.", + "expires_at": null + } + ] + }"#; + + let driver = Arc::new(MockDreamDriver { + response: valid_response.to_string(), + }); + + let extractions = vec![SessionExtraction { + preferences: vec!["dark mode".to_string()], + open_items: vec!["code review pending".to_string()], + ..SessionExtraction::default() + }]; + + let recent = vec![Message::user("Let's set up CI/CD next.")]; + let user_id = UserId::new(); + let config = CompactionConfig::default(); + + let result = dream( + driver, + "test-model", + &extractions, + &recent, + user_id, + &[("preferences".to_string(), "User prefers tea".to_string())], + &config, + ) + .await + .unwrap(); + + assert_eq!(result.topics.len(), 2); + assert_eq!(result.topics[0].topic, "preferences"); + assert_eq!(result.topics[1].topic, "open_items"); + assert_eq!(result.user_id, user_id); + assert!(result.superseded_topics.is_empty()); + } + + #[tokio::test] + async fn test_dream_conflict_resolution() { + let response_with_supersedes = r#"{ + "supersedes": ["old_preferences"], + "topics": [ + { + "topic": "preferences", + "summary": "Updated preferences", + "content": "User now prefers coffee (changed from tea).", + "expires_at": null + } + ] + }"#; + + let driver = Arc::new(MockDreamDriver { + response: response_with_supersedes.to_string(), + }); + + let user_id = UserId::new(); + let result = dream( + driver, + "test-model", + &[], + &[], + user_id, + &[( + "old_preferences".to_string(), + "User prefers tea".to_string(), + )], + &CompactionConfig::default(), + ) + .await + .unwrap(); + + assert_eq!( + result.superseded_topics, + vec!["old_preferences".to_string()] + ); + assert_eq!(result.topics.len(), 1); + assert_eq!(result.topics[0].topic, "preferences"); + } + + #[tokio::test] + async fn test_dream_expiry_tagging() { + let response_with_expiry = r#"{ + "supersedes": [], + "topics": [ + { + "topic": "upcoming_travel", + "summary": "Paris trip", + "content": "Flying to Paris on April 10.", + "expires_at": "2026-04-11T00:00:00Z" + }, + { + "topic": "preferences", + "summary": "Durable preferences", + "content": "Prefers concise answers.", + "expires_at": null + } + ] + }"#; + + let driver = Arc::new(MockDreamDriver { + response: response_with_expiry.to_string(), + }); + + let user_id = UserId::new(); + let result = dream( + driver, + "test-model", + &[], + &[], + user_id, + &[], + &CompactionConfig::default(), + ) + .await + .unwrap(); + + assert_eq!(result.topics.len(), 2); + assert!(result.topics[0].expires_at.is_some()); + assert!(result.topics[1].expires_at.is_none()); + } + + #[tokio::test] + async fn test_dream_fallback_on_parse_error() { + let driver = Arc::new(MockDreamDriver { + response: "this is not valid json at all!!!".to_string(), + }); + + let extractions = vec![SessionExtraction { + facts: vec!["User uses Rust".to_string()], + ..SessionExtraction::default() + }]; + + let config = CompactionConfig { + max_retries: 2, + ..CompactionConfig::default() + }; + + let user_id = UserId::new(); + let result = dream( + driver, + "test-model", + &extractions, + &[], + user_id, + &[], + &config, + ) + .await + .unwrap(); + + assert_eq!(result.topics.len(), 1); + assert_eq!(result.topics[0].topic, "general"); + assert!(result.topics[0].content.contains("User uses Rust")); + assert!(result.superseded_topics.is_empty()); + } +} diff --git a/crates/openfang-runtime/src/lib.rs b/crates/openfang-runtime/src/lib.rs index bde54ab199..23e25aca88 100644 --- a/crates/openfang-runtime/src/lib.rs +++ b/crates/openfang-runtime/src/lib.rs @@ -20,6 +20,7 @@ pub mod context_budget; pub mod context_overflow; pub mod copilot_oauth; pub mod docker_sandbox; +pub mod dreamer; pub mod drivers; pub mod embedding; pub mod graceful_shutdown; @@ -34,6 +35,7 @@ pub mod loop_guard; pub mod mcp; pub mod mcp_server; pub mod media_understanding; +pub mod mini_dream; pub mod model_catalog; pub mod process_manager; pub mod prompt_builder; diff --git a/crates/openfang-runtime/src/mini_dream.rs b/crates/openfang-runtime/src/mini_dream.rs new file mode 100644 index 0000000000..0b6db0bdbb --- /dev/null +++ b/crates/openfang-runtime/src/mini_dream.rs @@ -0,0 +1,192 @@ +//! Mini-dream: extract and persist facts from messages being trimmed during overflow recovery. +//! +//! When the context window overflows and old messages must be dropped, this runs a +//! dream pass on the messages about to be removed so their content is preserved in +//! user memory before they are discarded — replacing lossy truncation with semantic +//! preservation. +//! +//! Called by the agent loop immediately before `recover_from_overflow()` trims the +//! context. Non-fatal: errors are logged but never propagated. + +use crate::compactor::{extract_structured, CompactionConfig}; +use crate::dreamer::dream; +use crate::embedding::EmbeddingDriver; +use crate::llm_driver::LlmDriver; +use openfang_memory::user_memory::MemoryTopic; +use openfang_memory::MemorySubstrate; +use openfang_types::agent::UserId; +use openfang_types::message::Message; +use std::sync::Arc; +use tracing::{debug, info, warn}; + +/// Extract facts from `messages` (which are about to be trimmed) and persist them +/// to user memory. Non-fatal: errors are logged, never propagated. +/// +/// `embedding_driver` is optional — when provided, topic embeddings are stored for +/// semantic retrieval. +/// +/// Returns the number of memory topics written. +pub async fn run_mini_dream( + messages: &[Message], + user_id: UserId, + driver: Arc, + model: &str, + memory: &MemorySubstrate, + config: &CompactionConfig, + embedding_driver: Option<&(dyn EmbeddingDriver + Send + Sync)>, +) -> usize { + if messages.is_empty() { + return 0; + } + + debug!( + user_id = %user_id, + messages = messages.len(), + "Mini-dream: extracting facts from messages about to be trimmed" + ); + + // Step 1: Structured extraction from the messages being trimmed + let extraction = match extract_structured(driver.clone(), model, messages, None, config).await { + Ok(e) => e, + Err(e) => { + warn!("Mini-dream: structured extraction failed: {e}"); + return 0; + } + }; + + // Nothing worth persisting — skip the dream LLM call + if extraction.facts.is_empty() + && extraction.preferences.is_empty() + && extraction.decisions.is_empty() + && extraction.tasks.is_empty() + && extraction.open_items.is_empty() + { + debug!(user_id = %user_id, "Mini-dream: no facts extracted, skipping dream pass"); + return 0; + } + + // Step 2: Load existing topics for conflict detection / supersedes + let existing_topics: Vec<(String, String)> = memory + .user_topic_index(user_id) + .unwrap_or_default() + .into_iter() + .filter_map(|entry| { + memory + .user_topic(user_id, &entry.topic) + .ok() + .flatten() + .map(|t| (entry.topic, t.content)) + }) + .collect(); + + // Step 3: Dream pass — consolidate extraction into topic-organized memory + let result = match dream( + driver, + model, + &[extraction], + messages, + user_id, + &existing_topics, + config, + ) + .await + { + Ok(r) => r, + Err(e) => { + warn!("Mini-dream: dream pass failed: {e}"); + return 0; + } + }; + + // Step 4: Delete superseded topics, then write new ones + for superseded in &result.superseded_topics { + if let Err(e) = memory.delete_user_topic(user_id, superseded) { + warn!("Mini-dream: failed to delete superseded topic '{superseded}': {e}"); + } + } + + let now = chrono::Utc::now(); + let mut persisted = 0usize; + for dt in &result.topics { + let expires_at = dt + .expires_at + .as_deref() + .and_then(|s| s.parse::>().ok()); + let topic = MemoryTopic { + user_id, + topic: dt.topic.clone(), + summary: dt.summary.clone(), + content: dt.content.clone(), + updated_at: now, + expires_at, + }; + match memory.upsert_user_topic(&topic) { + Ok(()) => { + persisted += 1; + // Embed the topic content for semantic retrieval. + if let Some(emb) = embedding_driver { + match emb.embed_one(&dt.content).await { + Ok(vec) => { + if let Err(e) = + memory.store_user_topic_embedding(user_id, &dt.topic, &vec) + { + warn!( + "Mini-dream: failed to store embedding for '{}': {e}", + dt.topic + ); + } + } + Err(e) => warn!("Mini-dream: embedding failed for '{}': {e}", dt.topic), + } + } + } + Err(e) => warn!("Mini-dream: failed to persist topic '{}': {e}", dt.topic), + } + } + + info!( + user_id = %user_id, + topics = persisted, + superseded = result.superseded_topics.len(), + "Mini-dream: facts preserved from trimmed messages" + ); + persisted +} + +#[cfg(test)] +mod tests { + use openfang_types::agent::{AgentManifest, MemoryConfig, MemorySystem}; + + /// Behavioural gate: the agent loop must only call into `run_mini_dream` + /// when the agent's manifest opts in to structured memory. + /// + /// This test asserts the gate predicate directly — the call sites in + /// `agent_loop.rs` are `if manifest.memory.is_structured() { ... }`, + /// so a `false` here proves the extraction path is skipped, and a + /// `true` here proves it runs. + #[test] + fn test_structured_memory_gate_skips_default_agents() { + // Default agent — no [memory] block in TOML. + let default_agent = AgentManifest::default(); + assert!( + !default_agent.memory.is_structured(), + "default agents must NOT trigger structured extraction / dreamer" + ); + assert_eq!(default_agent.memory.system, MemorySystem::Summarization); + } + + #[test] + fn test_structured_memory_gate_fires_for_opted_in_agents() { + // Agent that opts in to structured memory. + let opted_in = AgentManifest { + memory: MemoryConfig { + system: MemorySystem::Structured, + }, + ..AgentManifest::default() + }; + assert!( + opted_in.memory.is_structured(), + "agents with memory.system = structured MUST trigger structured extraction / dreamer" + ); + } +} diff --git a/crates/openfang-types/src/config.rs b/crates/openfang-types/src/config.rs index a7e0fe23a7..835cff1c24 100644 --- a/crates/openfang-types/src/config.rs +++ b/crates/openfang-types/src/config.rs @@ -1305,6 +1305,39 @@ pub struct KernelConfig { /// ``` #[serde(default)] pub skills: HashMap>, + /// Session lifecycle configuration (inactivity timeouts → dream trigger). + #[serde(default)] + pub sessions: SessionsConfig, +} + +/// Session lifecycle configuration. +/// +/// Controls how long a session can sit idle before the kernel triggers +/// dream consolidation (structured-memory agents only). +/// +/// ```toml +/// [sessions] +/// gap_secs = 900 # inactivity → session end + dream (default 300) +/// email_gap_secs = 86400 # longer timeout for async channels +/// ``` +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default)] +pub struct SessionsConfig { + /// Seconds of inactivity before a session is considered ended and dream is triggered. + /// Default: 300 (5 minutes). Set to 0 to disable the dream lifecycle loop entirely. + pub gap_secs: u64, + /// Longer inactivity timeout for async channels like email. + /// Default: 86400 (24 hours). Reserved for future channel-aware gating. + pub email_gap_secs: u64, +} + +impl Default for SessionsConfig { + fn default() -> Self { + Self { + gap_secs: 300, + email_gap_secs: 86_400, + } + } } /// Heartbeat monitor settings exposed in `[heartbeat]` config section. @@ -1547,6 +1580,7 @@ impl Default for KernelConfig { workflows_dir: None, heartbeat: HeartbeatSettings::default(), skills: HashMap::new(), + sessions: SessionsConfig::default(), } } } diff --git a/crates/openfang-types/src/message.rs b/crates/openfang-types/src/message.rs index e798b7108c..4b8b324bb6 100644 --- a/crates/openfang-types/src/message.rs +++ b/crates/openfang-types/src/message.rs @@ -307,6 +307,22 @@ impl Message { } } + /// Create a context-injection message (tagged for exclusion from dream). + /// + /// Context-injection messages carry calendar / email / heartbeat summaries + /// that should be visible to the LLM at runtime but excluded from the + /// structured-memory extraction pipeline so they do not bleed into + /// long-term user memory. + pub fn context_injection(content: impl Into) -> Self { + Self { + msg_id: new_msg_id(), + provider_msg_id: None, + role: Role::User, + content: MessageContent::Text(content.into()), + source: Some(MessageSource::ContextInjection), + } + } + /// Create an assistant message with structured content blocks. /// /// Used to preserve `Thinking` blocks (with signatures and reasoning text) From e33a4c6c94015845597368d52c296cd04fdc4aab Mon Sep 17 00:00:00 2001 From: Philippe Branchu Date: Wed, 3 Jun 2026 07:11:34 +0000 Subject: [PATCH 05/12] fix(kernel): extract has_real_user_activity predicate + unit-test the gate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The dream activity gate (skip pure-tick / pure-context-injection sessions) was inlined in trigger_session_dream with zero direct test coverage. Pulled it out into a pub(crate) free function so the predicate can be unit-tested without spinning up a kernel, then added 8 tests covering the scenarios that matter for the thundering-herd fix: real user activity, autonomous ticks, scheduled ticks, mixed tick+real, pure context injections, empty sessions, assistant-only sessions, and ticks + context injections. This is the most critical test gap flagged by review — the predicate is the live protection against every heartbeat firing a dream pass. --- crates/openfang-kernel/src/kernel.rs | 163 ++++++++++++++++++++++++--- 1 file changed, 149 insertions(+), 14 deletions(-) diff --git a/crates/openfang-kernel/src/kernel.rs b/crates/openfang-kernel/src/kernel.rs index 29ae01a8c4..afc251834d 100644 --- a/crates/openfang-kernel/src/kernel.rs +++ b/crates/openfang-kernel/src/kernel.rs @@ -615,6 +615,34 @@ fn gethostname() -> Option { } } +/// Returns `true` if the session contains at least one message that represents +/// genuine user activity, as opposed to autonomous tick prompts or +/// context-injection summaries. +/// +/// A "real" message is: +/// - role = `User` +/// - **not** tagged `MessageSource::ContextInjection` (calendar/email/etc.) +/// - whose text does not start with `[AUTONOMOUS TICK]` or `[SCHEDULED TICK]` +/// +/// This is the thundering-herd guard for session dreaming: heartbeat-only +/// sessions (every agent fires a dream on every idle window) used to drown +/// the LLM in pointless extraction calls. Extracted to a free function so it +/// can be unit-tested directly without spinning up a kernel. +pub(crate) fn has_real_user_activity(messages: &[openfang_types::message::Message]) -> bool { + use openfang_types::message::{MessageSource, Role}; + + messages.iter().any(|msg| { + if msg.role != Role::User { + return false; + } + if msg.source == Some(MessageSource::ContextInjection) { + return false; + } + let text = msg.content.text_content(); + !text.starts_with("[AUTONOMOUS TICK]") && !text.starts_with("[SCHEDULED TICK]") + }) +} + impl OpenFangKernel { /// Boot the kernel with configuration from the given path. pub fn boot(config_path: Option<&Path>) -> KernelResult { @@ -5195,8 +5223,9 @@ impl OpenFangKernel { /// - The agent must opt in to structured memory (`manifest.memory.is_structured()`). /// - The session must contain *real* user activity — pure autonomous-tick sessions /// (where the only inputs are `[AUTONOMOUS TICK]` / `[SCHEDULED TICK]` prompts or - /// `ContextInjection` messages) are skipped to avoid the thundering-herd problem - /// that surfaced in production when every heartbeat fired a dream. + /// `ContextInjection` messages) are skipped via [`has_real_user_activity`] to + /// avoid the thundering-herd problem that surfaced in production when every + /// heartbeat fired a dream. pub async fn trigger_session_dream(self: &Arc, agent_id: AgentId) { use openfang_memory::user_memory::MemoryTopic; use openfang_runtime::compactor::{extract_structured, CompactionConfig}; @@ -5231,18 +5260,7 @@ impl OpenFangKernel { // tick prompts or context injections. A pure-tick session (heartbeat fires, // agent responds NO_REPLY, nothing else) is not worth consolidating. // This is the activity-gating fix from production (commit aa4ec5c on branchu). - let has_real_activity = session.messages.iter().any(|msg| { - if msg.role != openfang_types::message::Role::User { - return false; - } - if msg.source == Some(openfang_types::message::MessageSource::ContextInjection) { - return false; - } - let text = msg.content.text_content(); - !text.starts_with("[AUTONOMOUS TICK]") && !text.starts_with("[SCHEDULED TICK]") - }); - - if !has_real_activity { + if !has_real_user_activity(&session.messages) { debug!( agent_id = %agent_id, messages = session.messages.len(), @@ -10345,4 +10363,121 @@ system_prompt = "You are a test agent." kernel.shutdown(); } + + // ----------------------------------------------------------------- + // has_real_user_activity — the thundering-herd guard for dreaming. + // Pure-tick sessions must skip dream consolidation; sessions with any + // genuine user activity must proceed. + // ----------------------------------------------------------------- + + /// Scenario 1: real user activity → predicate true → dream would proceed. + #[test] + fn test_has_real_user_activity_with_user_message() { + let messages = vec![ + openfang_types::message::Message::user("Hello, how are you?"), + openfang_types::message::Message::assistant("I'm doing well, thanks!"), + ]; + assert!( + has_real_user_activity(&messages), + "session with a normal user message must count as real activity" + ); + } + + /// Scenario 2: pure tick-only session → predicate false → dream skipped. + #[test] + fn test_has_real_user_activity_only_autonomous_ticks() { + let messages = vec![ + openfang_types::message::Message::user("[AUTONOMOUS TICK] heartbeat"), + openfang_types::message::Message::assistant("NO_REPLY"), + openfang_types::message::Message::user("[AUTONOMOUS TICK] heartbeat"), + openfang_types::message::Message::assistant("NO_REPLY"), + ]; + assert!( + !has_real_user_activity(&messages), + "pure [AUTONOMOUS TICK] session must NOT count as real activity" + ); + } + + /// Pure scheduled-tick session is also skipped. + #[test] + fn test_has_real_user_activity_only_scheduled_ticks() { + let messages = vec![ + openfang_types::message::Message::user("[SCHEDULED TICK] daily summary"), + openfang_types::message::Message::assistant("NO_REPLY"), + ]; + assert!( + !has_real_user_activity(&messages), + "pure [SCHEDULED TICK] session must NOT count as real activity" + ); + } + + /// Scenario 3: mixed session (ticks + one real user message) → predicate true. + #[test] + fn test_has_real_user_activity_mixed_ticks_and_real() { + let messages = vec![ + openfang_types::message::Message::user("[AUTONOMOUS TICK] heartbeat"), + openfang_types::message::Message::assistant("NO_REPLY"), + openfang_types::message::Message::user("Hey, can you check my email?"), + openfang_types::message::Message::assistant("Sure, looking now."), + openfang_types::message::Message::user("[AUTONOMOUS TICK] heartbeat"), + ]; + assert!( + has_real_user_activity(&messages), + "mixed session with at least one real user message must count as real activity" + ); + } + + /// Scenario 4: pure context-injection session → predicate false. + #[test] + fn test_has_real_user_activity_only_context_injections() { + let messages = vec![ + openfang_types::message::Message::context_injection( + "Calendar: meeting with Alice at 3pm", + ), + openfang_types::message::Message::assistant("Noted."), + openfang_types::message::Message::context_injection("Email: 3 unread from Bob"), + ]; + assert!( + !has_real_user_activity(&messages), + "pure ContextInjection session must NOT count as real activity" + ); + } + + /// Empty message slice → predicate false (trivially no activity). + #[test] + fn test_has_real_user_activity_empty() { + let messages: Vec = vec![]; + assert!( + !has_real_user_activity(&messages), + "empty session must NOT count as real activity" + ); + } + + /// Only assistant messages → predicate false (assistant doesn't count). + #[test] + fn test_has_real_user_activity_only_assistant_messages() { + let messages = vec![ + openfang_types::message::Message::assistant("Hello!"), + openfang_types::message::Message::assistant("Anyone there?"), + ]; + assert!( + !has_real_user_activity(&messages), + "assistant-only session must NOT count as real activity" + ); + } + + /// Mixed ticks + context injections (no real user) → still false. + #[test] + fn test_has_real_user_activity_ticks_plus_context_injections() { + let messages = vec![ + openfang_types::message::Message::user("[AUTONOMOUS TICK] heartbeat"), + openfang_types::message::Message::context_injection("Calendar update"), + openfang_types::message::Message::user("[SCHEDULED TICK] morning"), + openfang_types::message::Message::assistant("NO_REPLY"), + ]; + assert!( + !has_real_user_activity(&messages), + "ticks + context-injections (no real user) must NOT count as real activity" + ); + } } From 683ea6087ef14996597316f500965493707239ee Mon Sep 17 00:00:00 2001 From: Philippe Branchu Date: Wed, 3 Jun 2026 07:13:19 +0000 Subject: [PATCH 06/12] fix(memory): infallible extract_structured + drop StubDriver fallback in dream MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit extract_structured was declared Result<_, String> but every code path returned Ok(fallback()) — the Err arm was unreachable. Callers in mini_dream.rs and the kernel dream entry point had dead match/unwrap_or arms that could never trigger. Made the signature infallible and stripped the dead error handling from both call sites. While touching the kernel call site, also fixed the StubDriver fallback that wrapped resolve_driver().ok().unwrap_or_else(...). StubDriver.complete() always errors, so feeding it to extract_structured burned the full retry budget (3 attempts) producing nothing but warn logs before returning the fallback. Resolve the driver up front and return early when missing — matching the pattern that already existed for the later dream() call in the same function. No behavior change for the happy path; eliminates wasted retries and noisy warns when no driver is configured, and removes dead error code. --- crates/openfang-kernel/src/kernel.rs | 48 ++++++++++------------- crates/openfang-runtime/src/compactor.rs | 12 +++--- crates/openfang-runtime/src/mini_dream.rs | 12 ++---- 3 files changed, 32 insertions(+), 40 deletions(-) diff --git a/crates/openfang-kernel/src/kernel.rs b/crates/openfang-kernel/src/kernel.rs index afc251834d..dd8c727c88 100644 --- a/crates/openfang-kernel/src/kernel.rs +++ b/crates/openfang-kernel/src/kernel.rs @@ -5283,32 +5283,34 @@ impl OpenFangKernel { let config = CompactionConfig::default(); let existing_extractions = self.memory.load_extractions(session_id).unwrap_or_default(); - // Always run a final extraction to capture anything since last compaction + // Resolve the driver up front. We do NOT fall back to StubDriver because + // StubDriver.complete() always errors, which burns the full retry budget + // inside extract_structured and emits noisy warns. If no driver is + // available, the dream pass is meaningless — return early to match the + // pattern used later in this function. + let driver = match self.resolve_driver(&entry.manifest) { + Ok(d) => d, + Err(e) => { + debug!(agent_id = %agent_id, "Dream: skipping — no LLM driver available: {e}"); + return; + } + }; + + // Always run a final extraction to capture anything since last compaction. + // `extract_structured` is infallible: internal failures fall back to the + // existing extraction (or default) and never error. let final_extraction = if !session.messages.is_empty() { let existing = existing_extractions.last().cloned(); - match extract_structured( - self.resolve_driver(&entry.manifest) - .ok() - .unwrap_or_else(|| { - Arc::new(crate::kernel::StubDriver) - as Arc - }), + let extraction = extract_structured( + driver.clone(), &entry.manifest.model.model, &session.messages, existing.as_ref(), &config, ) - .await - { - Ok(extraction) => { - let _ = self.memory.append_extraction(session_id, &extraction); - extraction - } - Err(e) => { - warn!(agent_id = %agent_id, "Dream: final extraction failed: {e}"); - existing_extractions.last().cloned().unwrap_or_default() - } - } + .await; + let _ = self.memory.append_extraction(session_id, &extraction); + extraction } else { existing_extractions.last().cloned().unwrap_or_default() }; @@ -5340,14 +5342,6 @@ impl OpenFangKernel { }) .collect(); - let driver = match self.resolve_driver(&entry.manifest) { - Ok(d) => d, - Err(e) => { - warn!(agent_id = %agent_id, "Dream: no driver available: {e}"); - return; - } - }; - // Take the last ~20 messages as the "recent tail" (not compacted) let recent_tail: Vec<_> = session .messages diff --git a/crates/openfang-runtime/src/compactor.rs b/crates/openfang-runtime/src/compactor.rs index 669143d590..3700365f57 100644 --- a/crates/openfang-runtime/src/compactor.rs +++ b/crates/openfang-runtime/src/compactor.rs @@ -130,13 +130,16 @@ pub fn count_tool_calls(messages: &[Message]) -> usize { /// On LLM error or persistent parse failure, returns the `existing` extraction /// (or `SessionExtraction::default()`) rather than propagating an error — the /// extraction path is non-critical and must never fail the agent turn. +/// +/// Signature is infallible by design: every failure mode falls back to the +/// existing extraction (or an empty one). Callers do not need error handling. pub async fn extract_structured( driver: Arc, model: &str, messages: &[Message], existing: Option<&SessionExtraction>, config: &CompactionConfig, -) -> Result { +) -> SessionExtraction { // Filter out context injection messages let filtered: Vec<&Message> = messages .iter() @@ -230,7 +233,7 @@ pub async fn extract_structured( match serde_json::from_str::(json_str) { Ok(extraction) => { info!("Structured extraction complete"); - return Ok(extraction); + return extraction; } Err(e) => { warn!(attempt, error = %e, "Failed to parse structured extraction JSON, retrying"); @@ -245,7 +248,7 @@ pub async fn extract_structured( // All retries failed — return existing or empty rather than erroring warn!("Structured extraction failed after all retries, returning fallback"); - Ok(fallback()) + fallback() } /// Estimate token count for a set of messages, optional system prompt, and tool definitions. @@ -1834,8 +1837,7 @@ mod tests { Some(&existing), &config, ) - .await - .unwrap(); + .await; // Returns the existing extraction as fallback assert_eq!(result.facts, vec!["User likes Rust".to_string()]); diff --git a/crates/openfang-runtime/src/mini_dream.rs b/crates/openfang-runtime/src/mini_dream.rs index 0b6db0bdbb..9068479346 100644 --- a/crates/openfang-runtime/src/mini_dream.rs +++ b/crates/openfang-runtime/src/mini_dream.rs @@ -45,14 +45,10 @@ pub async fn run_mini_dream( "Mini-dream: extracting facts from messages about to be trimmed" ); - // Step 1: Structured extraction from the messages being trimmed - let extraction = match extract_structured(driver.clone(), model, messages, None, config).await { - Ok(e) => e, - Err(e) => { - warn!("Mini-dream: structured extraction failed: {e}"); - return 0; - } - }; + // Step 1: Structured extraction from the messages being trimmed. + // `extract_structured` is infallible — internal LLM/parse failures fall back + // to an empty extraction, which the emptiness check below skips. + let extraction = extract_structured(driver.clone(), model, messages, None, config).await; // Nothing worth persisting — skip the dream LLM call if extraction.facts.is_empty() From 18f063545d34b31a6a1aee416281eb737eddee29 Mon Sep 17 00:00:00 2001 From: Philippe Branchu Date: Wed, 3 Jun 2026 07:14:24 +0000 Subject: [PATCH 07/12] fix(kernel): per-agent dream mutex to prevent concurrent dream tasks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The lifecycle loop fires a dream task per expired agent every 30s. If a dream takes longer than gap_secs while the same agent keeps receiving activity, the loop could re-enter and spawn a second dream task for the same agent — racing on extractions, embeddings, and user-topic writes. Added agent_dream_locks: DashMap>> on OpenFangKernel and gate every spawned dream task on `try_lock` of the per-agent mutex. `try_lock` (not `lock`) so iterations never queue — if the previous dream is still running, this tick logs a debug and skips, and the next 30s tick will check again. The mutex is separate from agent_msg_locks because that one serializes user turns vs user turns for the same agent; dreams should not block user turns and vice versa — only dream-vs-dream needs serialization. Added test_agent_dream_locks_serialize_per_agent covering: (a) second dream for same agent fails try_lock while first is in flight, (b) a different agent is not blocked, (c) once the first releases the next dispatch can lock again. --- crates/openfang-kernel/src/kernel.rs | 88 ++++++++++++++++++++++++++++ 1 file changed, 88 insertions(+) diff --git a/crates/openfang-kernel/src/kernel.rs b/crates/openfang-kernel/src/kernel.rs index dd8c727c88..42aefb0f99 100644 --- a/crates/openfang-kernel/src/kernel.rs +++ b/crates/openfang-kernel/src/kernel.rs @@ -182,6 +182,13 @@ pub struct OpenFangKernel { /// loop completion. The session lifecycle loop polls this to detect idle /// sessions and trigger dream consolidation for structured-memory agents. agent_last_active: dashmap::DashMap, + /// Per-agent in-flight dream mutex — guarantees only one dream task runs + /// per agent at a time. The lifecycle loop calls `try_lock` so iterations + /// never queue: if a previous dream is still running, this tick is skipped + /// and the next tick (30s later) will re-check. Separate from + /// `agent_msg_locks` because a dream and a user turn for the same agent + /// must not block each other — only dream-vs-dream is serialized. + agent_dream_locks: dashmap::DashMap>>, /// Weak self-reference for trigger dispatch (set after Arc wrapping). self_handle: OnceLock>, } @@ -1337,6 +1344,7 @@ impl OpenFangKernel { fallback_providers_override: std::sync::RwLock::new(None), agent_msg_locks: dashmap::DashMap::new(), agent_last_active: dashmap::DashMap::new(), + agent_dream_locks: dashmap::DashMap::new(), self_handle: OnceLock::new(), }; @@ -5207,7 +5215,23 @@ impl OpenFangKernel { for agent_id in expired { self.agent_last_active.remove(&agent_id); let kernel = Arc::clone(self); + // Per-agent in-flight dream mutex — if a previous dream task + // for this agent is still running, skip this iteration to + // avoid double-dispatch. We use `try_lock` (not `lock`) so + // tasks never queue: the next 30s tick will check again. + let dream_lock = self + .agent_dream_locks + .entry(agent_id) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) + .clone(); tokio::spawn(async move { + let Ok(_guard) = dream_lock.try_lock() else { + debug!( + agent_id = %agent_id, + "Dream: skipping — previous dream still in flight" + ); + return; + }; kernel.trigger_session_dream(agent_id).await; }); } @@ -10358,6 +10382,70 @@ system_prompt = "You are a test agent." kernel.shutdown(); } + // ----------------------------------------------------------------- + // agent_dream_locks — per-agent serialization of dream tasks. Two + // concurrent dispatches for the same agent must not both run. + // ----------------------------------------------------------------- + + /// Two concurrent dream tasks for the same agent must not run in parallel: + /// `try_lock` skips the second when the first is still holding the per-agent + /// dream mutex. We simulate this by acquiring the mutex (mirroring an + /// in-flight dream) and confirming a subsequent `try_lock` fails. A second + /// agent gets its own mutex, so its `try_lock` succeeds. + #[tokio::test] + async fn test_agent_dream_locks_serialize_per_agent() { + let tmp = tempfile::tempdir().unwrap(); + let kernel = minimal_kernel(&tmp); + + let agent_a = openfang_types::agent::AgentId::new(); + let agent_b = openfang_types::agent::AgentId::new(); + + // Simulate dream task A acquiring its lock. + let lock_a = kernel + .agent_dream_locks + .entry(agent_a) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) + .clone(); + let _guard_a = lock_a.lock().await; + + // A second dispatch for agent A should be skipped: try_lock fails. + let lock_a_again = kernel + .agent_dream_locks + .entry(agent_a) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) + .clone(); + assert!( + lock_a_again.try_lock().is_err(), + "concurrent dream for the same agent must be skipped (try_lock fails)" + ); + + // A dispatch for a different agent gets its own mutex — try_lock succeeds. + let lock_b = kernel + .agent_dream_locks + .entry(agent_b) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) + .clone(); + assert!( + lock_b.try_lock().is_ok(), + "different agent must NOT be blocked by another agent's in-flight dream" + ); + + drop(_guard_a); + + // Once A's lock is released, the next dispatch succeeds again. + let lock_a_after = kernel + .agent_dream_locks + .entry(agent_a) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) + .clone(); + assert!( + lock_a_after.try_lock().is_ok(), + "released lock must allow the next dream dispatch through" + ); + + kernel.shutdown(); + } + // ----------------------------------------------------------------- // has_real_user_activity — the thundering-herd guard for dreaming. // Pure-tick sessions must skip dream consolidation; sessions with any From 79d5a4fd3494f6ce33dfb2e9de5ffa7731fe7ecc Mon Sep 17 00:00:00 2001 From: Philippe Branchu Date: Wed, 3 Jun 2026 07:15:58 +0000 Subject: [PATCH 08/12] =?UTF-8?q?fix(memory):=20polish=20=E2=80=94=20drain?= =?UTF-8?q?=20count=20paired=20test,=20dead=5Fcode,=20visibility,=20race?= =?UTF-8?q?=20doc?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Five smaller cleanups bundled together: - context_overflow: added test_overflow_drain_count_matches_recover_from_overflow (stage 1 + stage 2 + below-threshold paired tests) so the two implementations cannot silently drift apart. overflow_drain_count mirrors stages 1+2 of recover_from_overflow by hand; if either threshold moves, the paired test catches it. - compactor: marked count_tool_calls + needs_extraction with #[allow(dead_code)] and a TODO(PR4-or-later) explaining the intent. Today the extraction path is triggered by the context-overflow signal (overflow_drain_count → mini-dream); these helpers are the alternative cadence-based gate that will land in a follow-up PR once agent_loop tracks per-session counters. Keeping them next to CompactionConfig so the policy stays co-located. - compactor: tightened build_conversation_text visibility from pub to pub(crate). It is only called inside this crate (compactor and dreamer). - kernel: added a NOTE in run_session_lifecycle_loop explaining the benign race between the agent_last_active snapshot and the per-agent remove. The race only delays a dream by one gap_secs window; the dream itself reads the live session, so content correctness is preserved. --- crates/openfang-kernel/src/kernel.rs | 6 ++ crates/openfang-runtime/src/compactor.rs | 15 ++- .../openfang-runtime/src/context_overflow.rs | 92 +++++++++++++++++++ 3 files changed, 112 insertions(+), 1 deletion(-) diff --git a/crates/openfang-kernel/src/kernel.rs b/crates/openfang-kernel/src/kernel.rs index 42aefb0f99..a3d9bf6e74 100644 --- a/crates/openfang-kernel/src/kernel.rs +++ b/crates/openfang-kernel/src/kernel.rs @@ -5204,6 +5204,12 @@ impl OpenFangKernel { break; } + // NOTE: Race window — if a new turn finishes between the snapshot + // taken here and the per-agent `remove` below, the fresh timestamp + // is wiped. Result is a delayed `gap_secs` window (dream fires at + // the next tick instead of being suppressed by the new activity). + // The dream task itself reads the live session, so the content is + // correct; only the timing is briefly off — benign. let now = std::time::Instant::now(); let expired: Vec = self .agent_last_active diff --git a/crates/openfang-runtime/src/compactor.rs b/crates/openfang-runtime/src/compactor.rs index 3700365f57..612b880c26 100644 --- a/crates/openfang-runtime/src/compactor.rs +++ b/crates/openfang-runtime/src/compactor.rs @@ -96,6 +96,15 @@ pub fn needs_compaction(session: &Session, config: &CompactionConfig) -> bool { /// Check whether a session needs structured extraction. /// Triggers on token count OR tool call count since last extraction. +/// +/// TODO(PR4-or-later): wire this in once `agent_loop` tracks per-session +/// `tokens_since_last_extraction` / `tool_calls_since_last_extraction` +/// counters. Today the extraction is gated by the context-overflow path +/// (`overflow_drain_count` → mini-dream), which is a different signal — +/// proximity to the context window, not extraction cadence. Keeping this +/// helper here so the policy lives next to `CompactionConfig`, but +/// `#[allow(dead_code)]` until the call site lands. +#[allow(dead_code)] pub fn needs_extraction( _messages_since_last: usize, tokens_since_last: usize, @@ -108,6 +117,10 @@ pub fn needs_extraction( } /// Count the number of tool calls (ToolUse blocks) in a slice of messages. +/// +/// TODO(PR4-or-later): wire into `agent_loop` alongside `needs_extraction` +/// once per-session counters are tracked. See `needs_extraction` above. +#[allow(dead_code)] pub fn count_tool_calls(messages: &[Message]) -> usize { messages .iter() @@ -491,7 +504,7 @@ fn is_oversized(message: &Message, config: &CompactionConfig) -> bool { /// /// Handles all content block types: text, tool use, tool result, image, unknown. /// Oversized messages are truncated inline with a marker. -pub fn build_conversation_text(messages: &[Message], config: &CompactionConfig) -> String { +pub(crate) fn build_conversation_text(messages: &[Message], config: &CompactionConfig) -> String { let mut conversation_text = String::new(); for msg in messages { diff --git a/crates/openfang-runtime/src/context_overflow.rs b/crates/openfang-runtime/src/context_overflow.rs index 6dcefd16fc..015f2a10e0 100644 --- a/crates/openfang-runtime/src/context_overflow.rs +++ b/crates/openfang-runtime/src/context_overflow.rs @@ -449,4 +449,96 @@ mod tests { let msgs = vec![Message::user("a"), Message::assistant("b")]; assert_eq!(safe_drain_boundary(&msgs, 2), 2); } + + /// Cross-check: `overflow_drain_count` must report exactly how many + /// leading messages `recover_from_overflow` will drain in stages 1+2. + /// If the thresholds in either function drift, this test catches it + /// before mini-dream silently sees a different message slice than + /// `recover_from_overflow` actually trims. + #[test] + fn test_overflow_drain_count_matches_recover_from_overflow_stage1() { + // Stage 1 zone: between 70% and 90% of context window. + // 1000-token context window: 70% = 700 tokens = 2800 chars, 90% = 2520 tokens. + // 20 messages of 150 chars each ≈ 3000+ chars total → past 70%. + let mut msgs_for_recover = make_messages(20, 150); + let msgs_for_count = msgs_for_recover.clone(); + let system_prompt = "system"; + let tools: Vec = vec![]; + let context_window = 1000; + + let predicted = + overflow_drain_count(&msgs_for_count, system_prompt, &tools, context_window); + let original_len = msgs_for_recover.len(); + let stage = + recover_from_overflow(&mut msgs_for_recover, system_prompt, &tools, context_window); + + if let RecoveryStage::AutoCompaction { removed } = stage { + assert_eq!( + predicted, removed, + "overflow_drain_count must match Stage 1 (AutoCompaction) drain" + ); + assert_eq!( + original_len - msgs_for_recover.len(), + predicted, + "drain count must equal predicted" + ); + } else if matches!(stage, RecoveryStage::OverflowCompaction { .. }) { + // Cascaded to stage 2 — predicted should still match the leading + // drain count, which equals the stage 2 drain since Stage 1 drained + // some, then Stage 2 drained the rest. We accept this as long as + // predicted is non-zero. + assert!( + predicted > 0, + "predicted must be non-zero for overflowing input" + ); + } + } + + #[test] + fn test_overflow_drain_count_matches_recover_from_overflow_stage2() { + // Stage 2 zone: > 90%. 30 messages of 200 chars ≈ 6000 chars in a + // 1000-token (4000-char) window → way past 90%. + let mut msgs_for_recover = make_messages(30, 200); + let msgs_for_count = msgs_for_recover.clone(); + let system_prompt = "system"; + let tools: Vec = vec![]; + let context_window = 1000; + + let predicted = + overflow_drain_count(&msgs_for_count, system_prompt, &tools, context_window); + let stage = + recover_from_overflow(&mut msgs_for_recover, system_prompt, &tools, context_window); + + // Whichever stage fires, the prediction must equal what was drained. + let actual_drained = match stage { + RecoveryStage::AutoCompaction { removed } => removed, + RecoveryStage::OverflowCompaction { removed } => removed, + RecoveryStage::ToolResultTruncation { .. } | RecoveryStage::FinalError => { + // Stage 2 still ran first; check that predicted matches what + // it would have drained. + assert!(predicted > 0, "predicted must be non-zero past 90%"); + return; + } + RecoveryStage::None => { + panic!("expected recovery to fire"); + } + }; + + assert_eq!( + predicted, actual_drained, + "overflow_drain_count must match recover_from_overflow drain" + ); + } + + #[test] + fn test_overflow_drain_count_zero_below_threshold() { + // Below 70% — both functions agree: no drain, no recovery. + let msgs = make_messages(2, 100); + let predicted = overflow_drain_count(&msgs, "sys", &[], 200_000); + assert_eq!(predicted, 0, "below threshold must report 0 drain"); + + let mut msgs_mut = msgs; + let stage = recover_from_overflow(&mut msgs_mut, "sys", &[], 200_000); + assert_eq!(stage, RecoveryStage::None); + } } From cc0aa33052fe630b6eb639d142b9c695ddf62909 Mon Sep 17 00:00:00 2001 From: Philippe Branchu Date: Wed, 3 Jun 2026 07:40:08 +0000 Subject: [PATCH 09/12] =?UTF-8?q?feat(memory):=20dashboard=20UI=20?= =?UTF-8?q?=E2=80=94=20opt-in=20selector=20+=20memory=20management?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Dashboard surfaces for the structured memory feature. Per-agent Memory System dropdown in the agent Config tab (opt-in). New Users page with Memory and Extraction Audit tabs for viewing/deleting/exporting per-user accumulated memory. Pure UI work — all backend routes and the PATCH field already shipped in the memory storage PR. - Memory System selector in the agent detail Config tab: -