From a2ced40ef9c272b2ba43c0744c2bd17373afca70 Mon Sep 17 00:00:00 2001 From: YellowSnnowmann <167776381+YellowSnnowmann@users.noreply.github.com> Date: Tue, 12 May 2026 00:37:59 +0530 Subject: [PATCH] feat(memory): add tool-scoped memory with durable rule capture, RPC APIs, and prompt pinning (#1487) Co-authored-by: Claude Sonnet 4.6 --- gitbooks/SUMMARY.md | 1 + gitbooks/features/native-tools/tool-memory.md | 82 +++ src/openhuman/about_app/catalog.rs | 14 + .../agent/harness/session/builder.rs | 78 +++ src/openhuman/agent/prompts/mod.rs | 35 ++ src/openhuman/config/schema/learning.rs | 13 + src/openhuman/config/schema/load.rs | 8 + src/openhuman/memory/mod.rs | 6 + src/openhuman/memory/ops/mod.rs | 6 + src/openhuman/memory/ops/tool_memory.rs | 174 +++++++ src/openhuman/memory/schemas/mod.rs | 6 + src/openhuman/memory/schemas/tool_memory.rs | 269 ++++++++++ src/openhuman/memory/schemas_tests.rs | 7 + src/openhuman/memory/store/client.rs | 9 + src/openhuman/memory/tool_memory/capture.rs | 481 ++++++++++++++++++ src/openhuman/memory/tool_memory/mod.rs | 44 ++ src/openhuman/memory/tool_memory/prompt.rs | 250 +++++++++ src/openhuman/memory/tool_memory/store.rs | 269 ++++++++++ .../memory/tool_memory/store_tests.rs | 375 ++++++++++++++ .../memory/tool_memory/test_helpers.rs | 107 ++++ src/openhuman/memory/tool_memory/types.rs | 239 +++++++++ 21 files changed, 2473 insertions(+) create mode 100644 gitbooks/features/native-tools/tool-memory.md create mode 100644 src/openhuman/memory/ops/tool_memory.rs create mode 100644 src/openhuman/memory/schemas/tool_memory.rs create mode 100644 src/openhuman/memory/tool_memory/capture.rs create mode 100644 src/openhuman/memory/tool_memory/mod.rs create mode 100644 src/openhuman/memory/tool_memory/prompt.rs create mode 100644 src/openhuman/memory/tool_memory/store.rs create mode 100644 src/openhuman/memory/tool_memory/store_tests.rs create mode 100644 src/openhuman/memory/tool_memory/test_helpers.rs create mode 100644 src/openhuman/memory/tool_memory/types.rs diff --git a/gitbooks/SUMMARY.md b/gitbooks/SUMMARY.md index 2c66d0ae2..8e1613f91 100644 --- a/gitbooks/SUMMARY.md +++ b/gitbooks/SUMMARY.md @@ -25,6 +25,7 @@ * [Cron & Scheduling](features/native-tools/cron.md) * [Voice](features/native-tools/voice.md) * [Memory Tools](features/native-tools/memory-tools.md) + * [Tool-Scoped Memory](features/native-tools/tool-memory.md) * [Third-party Integrations](features/native-tools/integrations.md) * [Agent Coordination](features/native-tools/agent-coordination.md) * [System & Utilities](features/native-tools/system-and-utilities.md) diff --git a/gitbooks/features/native-tools/tool-memory.md b/gitbooks/features/native-tools/tool-memory.md new file mode 100644 index 000000000..291b91cda --- /dev/null +++ b/gitbooks/features/native-tools/tool-memory.md @@ -0,0 +1,82 @@ +--- +description: Durable, tool-scoped rules for safety-critical guidance and learnings. +icon: shield-check +--- + +# Tool-Scoped Memory + +The tool-scoped memory layer captures **actionable guidance** about how the agent should use specific tools - separate from the [Memory Tools](memory-tools.md) general-purpose recall and from the `tool_effectiveness` statistics namespace. It is the surface that turns "never email Sarah" into a hard constraint the agent has to obey on every subsequent turn. + +It implements [issue #1400](https://github.com/tinyhumansai/openhuman/issues/1400) - a first-class storage and retrieval system for durable learnings and high-priority rules. + +## What it stores + +Every tool gets its own namespace, **`tool-{tool_name}`**, distinct from `global`, `skill-{id}`, and the statistics-only `tool_effectiveness` namespace. Inside it, each entry is a `ToolMemoryRule`: + +| Field | Purpose | +| ------------ | ------- | +| `id` | Stable per-rule UUID. Upserts replay the same id. | +| `tool_name` | The tool the rule applies to (e.g. `send_email`, `shell`). | +| `rule` | Natural-language guidance the agent must follow. | +| `priority` | `critical`, `high`, or `normal`. Drives retrieval + compression. | +| `source` | `user_explicit`, `post_turn`, or `programmatic` - provenance. | +| `tags` | Free-form labels (`safety`, `permission`, ...). | +| `created_at` / `updated_at` | RFC3339 timestamps. | + +Statistics (`tool_effectiveness/tool/{name}`) and rules (`tool-{name}/rule/{id}`) live in *different* namespaces by design - one tracks "what happened", the other tracks "what to do about it". + +## Priority levels + +| Priority | Where it lives | Compression-resistant? | +| ---------- | -------------- | ---------------------- | +| `critical` | Pinned into the **system prompt** via `ToolMemoryRulesSection`. | **Yes** - the system prompt is frozen per-session and never rewritten by the mid-session compactor. | +| `high` | Same system prompt block, ranked below critical. | **Yes** - same mechanism. | +| `normal` | Stored in the namespace; retrieved on demand via `memory_recall`. | No - eligible for compression like any other namespaced memory. | + +The compression-resistance property is structural: critical and high rules ride in the *system prompt*, which the inference backend's prefix cache keeps frozen for the entire session. There is no way for token compression to silently drop a `critical` rule. + +## Capture pipeline + +Two automatic capture paths fire after every turn (via `ToolMemoryCaptureHook`): + +1. **User edicts** - sentences like `never `, `don't ...`, `do not ...`, or `stop ing ...` in the user message are promoted to a **Critical** rule on the matching tool. Common-noun aliases map `"email"` to a tool named `send_email`, `"shell"` to `bash`/`exec`, etc.; when no alias matches, the rule lands on the first tool that ran in the turn so it stays adjacent to the relevant call site. +2. **Repeated tool failures** - a tool that fails twice or more in a single turn earns a **Normal**-priority observation, with the failure class summarized inline so the agent has context next time it considers that tool. + +The hook is enabled by default whenever the learning subsystem is on. Disable selectively with `OPENHUMAN_LEARNING_TOOL_MEMORY_CAPTURE_ENABLED=0`. + +## Retrieval at tool-selection time + +At session start the harness pre-fetches every Critical and High rule via `ToolMemoryStore::rules_for_prompt`, renders them into the `## Tool-scoped rules` block, and pins the block into the system prompt. Because the prompt is frozen for the lifetime of the session, the rules are visible on every turn at tool-selection time and before any actual tool execution. + +Lower-priority guidance stays out of the prompt budget; the agent reaches it on demand by calling `memory_recall` against the `tool-{name}` namespace. + +## RPC surface + +Six methods are exposed under the `memory` namespace: + +| Method | Purpose | +| ----------------------------------- | ------- | +| `memory.tool_rule_put` | Upsert a rule. Use `priority='critical'` for safety-critical entries. | +| `memory.tool_rule_get` | Fetch a rule by `(tool_name, id)`. | +| `memory.tool_rule_list` | List all rules for a tool, sorted by priority + freshness. | +| `memory.tool_rule_delete` | Delete a rule. | +| `memory.tool_rules_for_prompt` | Return the rendered Markdown block + structured snapshot - what the session builder pins. | +| `memory.tool_rules_json` | Raw JSON list (for envelope consumers). | + +JSON payloads use snake_case (`priority: "critical"`, `source: "user_explicit"`). Every method goes through the same `active_memory_client` plumbing as the rest of the memory RPCs. + +## End-to-end safety case + +The "never email Sarah" path is covered as a regression test: + +1. User says *"Never email Sarah at sarah@example.com."* during a turn that called `send_email`. +2. `ToolMemoryCaptureHook` extracts the edict, maps the `email` alias to the `send_email` tool, and writes a Critical rule under `tool-send_email/rule/{uuid}`. +3. On the next session, `prefetch_tool_memory_rules_blocking` pulls every Critical and High rule and the session builder appends a `ToolMemoryRulesSection` to the system prompt. +4. The agent sees `### \`send_email\`` followed by `- **[critical]** Never email Sarah at sarah@example.com.` before ever choosing a tool, and the rule survives any mid-session token compression. + +Coverage and the integration test live in `src/openhuman/memory/tool_memory/`. + +## See also + +- [Memory Tools](memory-tools.md) - general-purpose `recall`, `store`, `forget`. +- [Smart Token Compression](../token-compression.md) - what the system prompt is protected against. diff --git a/src/openhuman/about_app/catalog.rs b/src/openhuman/about_app/catalog.rs index c7bd72929..94e3fc70c 100644 --- a/src/openhuman/about_app/catalog.rs +++ b/src/openhuman/about_app/catalog.rs @@ -206,6 +206,20 @@ const CAPABILITIES: &[Capability] = &[ status: CapabilityStatus::Beta, privacy: None, }, + Capability { + id: "intelligence.tool_scoped_memory", + name: "Tool-Scoped Memory Rules", + domain: "intelligence", + category: CapabilityCategory::Intelligence, + description: "Store durable, tool-specific rules and corrections that survive context \ + compression. Critical-priority rules (e.g. 'never email Sarah') are pinned into the \ + system prompt at session start. Captured automatically from user edicts and repeated \ + tool failures; also writable programmatically via the memory.tool_rule_* RPC surface.", + how_to: "Automatic — user edicts are captured after every turn. Manage via \ + memory.tool_rule_put / memory.tool_rule_list / memory.tool_rule_delete (RPC).", + status: CapabilityStatus::Beta, + privacy: LOCAL_RAW, + }, Capability { id: "intelligence.memory_tree_retrieval", name: "Memory Tree Retrieval (chat)", diff --git a/src/openhuman/agent/harness/session/builder.rs b/src/openhuman/agent/harness/session/builder.rs index e3c20fe9b..d532acf0b 100644 --- a/src/openhuman/agent/harness/session/builder.rs +++ b/src/openhuman/agent/harness/session/builder.rs @@ -840,6 +840,13 @@ impl Agent { ))); log::info!("[learning] tool_tracker hook registered"); } + + if config.learning.tool_memory_capture_enabled { + post_turn_hooks.push(Arc::new( + crate::openhuman::memory::ToolMemoryCaptureHook::new(memory.clone(), true), + )); + log::info!("[learning] tool_memory_capture hook registered"); + } } // Resolve the per-agent delegation tool set and visible-tool @@ -970,6 +977,27 @@ impl Agent { .filter(|t| !existing_names.contains(t.name())), ); + // Pre-fetch Critical + High priority tool-scoped memory rules so they + // pin into the (compression-resistant) system prompt for the whole + // session. Done here — after the tool list is finalised — so we only + // fetch rules for tools this agent can actually use. Skipped when + // `learning.enabled` is false (no new rules are written in that mode, + // and users who opt out of learning expect no stored rules to surface) + // or when the runtime cannot host a synchronous bridge (single-threaded + // test harnesses). + if config.learning.enabled && config.learning.tool_memory_capture_enabled { + let agent_tool_names: Vec = + tools.iter().map(|t| t.name().to_string()).collect(); + let pinned = prefetch_tool_memory_rules_blocking(memory.clone(), &agent_tool_names); + if !pinned.is_empty() { + log::info!( + "[memory::tool_memory] pinning {} tool-scoped rule(s) into system prompt", + pinned.len() + ); + prompt_builder = prompt_builder.with_tool_memory_rules(pinned); + } + } + // Build the P-Format registry AFTER the tool list is finalised // (including orchestrator tools) so every tool gets a signature // entry. The registry is self-contained — it doesn't hold a @@ -1152,3 +1180,53 @@ impl Agent { builder.build() } } + +/// (#1400) Best-effort synchronous prefetch of eager tool-scoped rules. +/// +/// `from_config_*` is sync but typically runs inside a multi-threaded +/// Tokio runtime (the agent harness path from the channels runtime). +/// We use `block_in_place` + the current runtime handle to call the +/// async store API without restructuring the whole session builder. +/// +/// Returns an empty `Vec` (rather than erroring) when: +/// - no Tokio runtime is active (e.g. a sync CLI bootstrap), +/// - the runtime is single-threaded (`block_in_place` would panic), +/// - or the underlying `rules_for_prompt` call returns an error +/// (e.g. the memory backend isn't ready yet). +/// +/// Critical / High rules captured later in the session are still +/// available via the `memory_tool_rules_for_prompt` RPC; this prefetch +/// merely seeds the rules that exist at session start. +fn prefetch_tool_memory_rules_blocking( + memory: Arc, + tool_names: &[String], +) -> Vec { + let Ok(handle) = tokio::runtime::Handle::try_current() else { + return Vec::new(); + }; + if handle.runtime_flavor() != tokio::runtime::RuntimeFlavor::MultiThread { + return Vec::new(); + } + let tool_names = tool_names.to_vec(); + tokio::task::block_in_place(|| { + handle.block_on(async move { + let store = crate::openhuman::memory::ToolMemoryStore::new(memory); + match store.rules_for_prompt(&tool_names).await { + Ok(grouped) => { + let mut flat: Vec<_> = grouped.into_values().flatten().collect(); + flat.sort_by(|a, b| { + b.priority + .cmp(&a.priority) + .then_with(|| a.tool_name.cmp(&b.tool_name)) + .then_with(|| a.rule.cmp(&b.rule)) + }); + flat + } + Err(err) => { + log::warn!("[memory::tool_memory] prefetch failed: {err}"); + Vec::new() + } + } + }) + }) +} diff --git a/src/openhuman/agent/prompts/mod.rs b/src/openhuman/agent/prompts/mod.rs index 3c8e4f966..561f25f2f 100644 --- a/src/openhuman/agent/prompts/mod.rs +++ b/src/openhuman/agent/prompts/mod.rs @@ -170,6 +170,41 @@ impl SystemPromptBuilder { self } + /// Append a [`ToolMemoryRulesSection`] carrying a pre-fetched + /// snapshot of Critical / High priority tool-scoped rules (#1400). + /// + /// Snapshot semantics — the rules are baked into the section at + /// construction so the rendered system prompt stays byte-identical + /// for the lifetime of the session. The session builder is + /// responsible for pre-fetching via + /// [`crate::openhuman::memory::ToolMemoryStore::rules_for_prompt`] + /// (or the `memory_tool_rules_for_prompt` RPC) before invoking + /// this method. + /// + /// No-op when `rules` is empty. + pub fn with_tool_memory_rules( + mut self, + rules: Vec, + ) -> Self { + if rules.is_empty() { + return self; + } + // Insert before the tool-catalogue section so these rules appear + // adjacent to the tool listings and survive tail-biased trimming. + // Falls back to push when no tools section is present. + let section: Box = + Box::new(crate::openhuman::memory::ToolMemoryRulesSection::new(rules)); + let tools_idx = self + .sections + .iter() + .position(|s| s.name() == "tools" || s.name() == "tool_catalogue"); + match tools_idx { + Some(idx) => self.sections.insert(idx, section), + None => self.sections.push(section), + } + self + } + /// Append a "Memory context" section carrying the resolved chunks the /// subconscious LLM cited when it produced the reflection that /// spawned this thread (#623). diff --git a/src/openhuman/config/schema/learning.rs b/src/openhuman/config/schema/learning.rs index 0d4f28689..c85d65449 100644 --- a/src/openhuman/config/schema/learning.rs +++ b/src/openhuman/config/schema/learning.rs @@ -35,6 +35,18 @@ pub struct LearningConfig { #[serde(default = "default_true")] pub tool_tracking_enabled: bool, + /// Enable the tool-scoped memory capture hook (see + /// [`crate::openhuman::memory::tool_memory::ToolMemoryCaptureHook`]). + /// + /// When enabled, the hook records user edicts ("never email Sarah") + /// as `Critical`-priority rules in the `tool-{name}` memory + /// namespace, and tallies repeated tool failures into + /// `Normal`-priority observations. Defaults to true when learning + /// is enabled — set to false to disable durable rule capture + /// without turning off learning entirely. + #[serde(default = "default_true")] + pub tool_memory_capture_enabled: bool, + /// Which LLM to use for reflection. Default: local (Ollama). #[serde(default)] pub reflection_source: ReflectionSource, @@ -91,6 +103,7 @@ impl Default for LearningConfig { reflection_enabled: default_true(), user_profile_enabled: default_true(), tool_tracking_enabled: default_true(), + tool_memory_capture_enabled: default_true(), reflection_source: ReflectionSource::default(), max_reflections_per_session: default_max_reflections(), min_turn_complexity: default_min_turn_complexity(), diff --git a/src/openhuman/config/schema/load.rs b/src/openhuman/config/schema/load.rs index 30de6da5d..b08c29fe9 100644 --- a/src/openhuman/config/schema/load.rs +++ b/src/openhuman/config/schema/load.rs @@ -898,6 +898,14 @@ impl Config { _ => {} } } + if let Some(flag) = env.get("OPENHUMAN_LEARNING_TOOL_MEMORY_CAPTURE_ENABLED") { + if let Some(enabled) = parse_env_bool( + "OPENHUMAN_LEARNING_TOOL_MEMORY_CAPTURE_ENABLED", + flag.as_str(), + ) { + self.learning.tool_memory_capture_enabled = enabled; + } + } if let Some(source) = env.get("OPENHUMAN_LEARNING_REFLECTION_SOURCE") { let normalized = source.trim().to_ascii_lowercase(); match normalized.as_str() { diff --git a/src/openhuman/memory/mod.rs b/src/openhuman/memory/mod.rs index c6a202a4f..e860eca7a 100644 --- a/src/openhuman/memory/mod.rs +++ b/src/openhuman/memory/mod.rs @@ -15,6 +15,7 @@ pub mod safety; pub mod schemas; pub mod store; pub mod sync_status; +pub mod tool_memory; pub mod traits; pub mod tree; @@ -40,6 +41,11 @@ pub use sync_status::{ all_memory_sync_status_controller_schemas, all_memory_sync_status_registered_controllers, FreshnessLabel, MemorySyncStatus, }; +pub use tool_memory::{ + render_tool_memory_rules, tool_memory_namespace, ToolMemoryCaptureHook, ToolMemoryPriority, + ToolMemoryRule, ToolMemoryRulesSection, ToolMemorySource, ToolMemoryStore, TOOL_MEMORY_HEADING, + TOOL_MEMORY_PROMPT_CAP, +}; pub use traits::{Memory, MemoryCategory, MemoryEntry, NamespaceSummary, RecallOpts}; pub use tree::{ all_memory_tree_controller_schemas, all_memory_tree_registered_controllers, diff --git a/src/openhuman/memory/ops/mod.rs b/src/openhuman/memory/ops/mod.rs index 55c17999f..3676e4107 100644 --- a/src/openhuman/memory/ops/mod.rs +++ b/src/openhuman/memory/ops/mod.rs @@ -26,6 +26,7 @@ pub mod helpers; pub mod kv_graph; pub mod learn; pub mod sync; +pub mod tool_memory; // --------------------------------------------------------------------------- // Re-exports preserving the previous flat `memory::ops::*` surface. @@ -48,6 +49,11 @@ pub use sync::{ memory_ingestion_status, memory_sync_all, memory_sync_channel, IngestionStatusResult, SyncAllResult, SyncChannelParams, SyncChannelResult, }; +pub use tool_memory::{ + tool_rule_delete, tool_rule_get, tool_rule_list, tool_rule_put, tool_rules_for_prompt, + tool_rules_json, ToolRuleListParams, ToolRulePutParams, ToolRuleRefParams, + ToolRulesForPromptParams, ToolRulesForPromptResult, +}; // --------------------------------------------------------------------------- // Test-only re-exports — keep the existing `ops_tests.rs` happy without diff --git a/src/openhuman/memory/ops/tool_memory.rs b/src/openhuman/memory/ops/tool_memory.rs new file mode 100644 index 000000000..495a2b0b9 --- /dev/null +++ b/src/openhuman/memory/ops/tool_memory.rs @@ -0,0 +1,174 @@ +//! RPC handlers for the tool-scoped memory layer (see +//! [`crate::openhuman::memory::tool_memory`]). +//! +//! All handlers go through [`active_memory_client`] so they hit the +//! same `UnifiedMemory` backend the rest of the memory RPCs use, and +//! the namespace they touch is exactly `tool-{tool_name}` — never +//! `global` or `tool_effectiveness`. + +use serde::Deserialize; +use serde_json::Value; + +use crate::openhuman::memory::tool_memory::{ + ToolMemoryPriority, ToolMemoryRule, ToolMemorySource, ToolMemoryStore, +}; +use crate::rpc::RpcOutcome; + +use super::helpers::active_memory_client; + +/// Parameters for `memory_tool_rule_put`. +#[derive(Debug, Deserialize)] +pub struct ToolRulePutParams { + /// Tool the rule applies to (e.g. `email`, `shell`). + pub tool_name: String, + /// Natural-language rule body. + pub rule: String, + /// Priority/criticality. Defaults to `normal` when omitted. + #[serde(default)] + pub priority: Option, + /// Provenance — defaults to `programmatic` when omitted. + #[serde(default)] + pub source: Option, + /// Optional tags for filtering. + #[serde(default)] + pub tags: Vec, + /// Optional rule id — when supplied, the call upserts in place + /// rather than creating a new entry. + #[serde(default)] + pub id: Option, +} + +/// Parameters for `memory_tool_rule_get` / `memory_tool_rule_delete`. +#[derive(Debug, Deserialize)] +pub struct ToolRuleRefParams { + pub tool_name: String, + pub id: String, +} + +/// Parameters for `memory_tool_rule_list`. +#[derive(Debug, Deserialize)] +pub struct ToolRuleListParams { + pub tool_name: String, +} + +/// Parameters for `memory_tool_rules_for_prompt`. +#[derive(Debug, Deserialize, Default)] +pub struct ToolRulesForPromptParams { + /// Constrain the result to these tools. Empty (or omitted) scans + /// every known tool namespace. + #[serde(default)] + pub tools: Vec, +} + +async fn open_store() -> Result { + let client = active_memory_client().await?; + Ok(ToolMemoryStore::new(client.memory_handle())) +} + +/// Upsert a tool-scoped memory rule. +pub async fn tool_rule_put( + params: ToolRulePutParams, +) -> Result, String> { + log::debug!("[tool-memory] rpc tool_rule_put tool={}", params.tool_name); + let store = open_store().await?; + let mut rule = ToolMemoryRule::new( + ¶ms.tool_name, + ¶ms.rule, + params.priority.unwrap_or_default(), + params.source.unwrap_or_default(), + ); + rule.tags = params.tags; + if let Some(id) = params.id { + if !id.trim().is_empty() { + rule.id = id; + } + } + let stored = store.put_rule(rule).await?; + Ok(RpcOutcome::single_log(stored, "tool memory rule stored")) +} + +/// Fetch a tool-scoped rule by id. +pub async fn tool_rule_get( + params: ToolRuleRefParams, +) -> Result>, String> { + log::debug!( + "[tool-memory] rpc tool_rule_get tool={} id={}", + params.tool_name, + params.id + ); + let store = open_store().await?; + let rule = store.get_rule(¶ms.tool_name, ¶ms.id).await?; + Ok(RpcOutcome::single_log(rule, "tool memory rule fetched")) +} + +/// List every tool-scoped rule for a tool. +pub async fn tool_rule_list( + params: ToolRuleListParams, +) -> Result>, String> { + log::debug!("[tool-memory] rpc tool_rule_list tool={}", params.tool_name); + let store = open_store().await?; + let rules = store.list_rules(¶ms.tool_name).await?; + Ok(RpcOutcome::single_log(rules, "tool memory rules listed")) +} + +/// Delete a tool-scoped rule by id. +pub async fn tool_rule_delete(params: ToolRuleRefParams) -> Result, String> { + log::debug!( + "[tool-memory] rpc tool_rule_delete tool={} id={}", + params.tool_name, + params.id + ); + let store = open_store().await?; + let deleted = store.delete_rule(¶ms.tool_name, ¶ms.id).await?; + Ok(RpcOutcome::single_log(deleted, "tool memory rule deleted")) +} + +/// Return the rendered prompt block plus the structured rule list for +/// the caller-supplied set of tools. Used by the session builder to +/// pin Critical / High rules into the system prompt. +#[derive(Debug, serde::Serialize)] +pub struct ToolRulesForPromptResult { + /// Pre-rendered Markdown block, ready for injection. + pub rendered: String, + /// Underlying rule snapshot the renderer used. + pub rules: Vec, +} + +/// Pre-fetch Critical + High priority rules for prompt injection. +pub async fn tool_rules_for_prompt( + params: ToolRulesForPromptParams, +) -> Result, String> { + log::debug!( + "[tool-memory] rpc tool_rules_for_prompt tools={:?}", + params.tools + ); + let store = open_store().await?; + let grouped = store.rules_for_prompt(¶ms.tools).await?; + let mut flat: Vec = grouped.into_values().flatten().collect(); + flat.sort_by(|a, b| { + b.priority + .cmp(&a.priority) + .then_with(|| a.tool_name.cmp(&b.tool_name)) + .then_with(|| a.rule.cmp(&b.rule)) + }); + let rendered = crate::openhuman::memory::tool_memory::render_tool_memory_rules(&flat); + Ok(RpcOutcome::single_log( + ToolRulesForPromptResult { + rendered, + rules: flat, + }, + "tool memory rules prepared for prompt", + )) +} + +/// Render the raw JSON form of a tool's rules, useful for envelope +/// consumers that want the unfiltered list. +pub async fn tool_rules_json(params: ToolRuleListParams) -> Result, String> { + log::debug!( + "[tool-memory] rpc tool_rules_json tool={}", + params.tool_name + ); + let store = open_store().await?; + let value = store.list_rules_json(¶ms.tool_name).await?; + Ok(RpcOutcome::single_log(value, "tool memory rules json")) +} diff --git a/src/openhuman/memory/schemas/mod.rs b/src/openhuman/memory/schemas/mod.rs index e7c775b4b..de11f2de5 100644 --- a/src/openhuman/memory/schemas/mod.rs +++ b/src/openhuman/memory/schemas/mod.rs @@ -25,6 +25,7 @@ mod files; mod kv_graph; mod learn; mod sync; +mod tool_memory; // --------------------------------------------------------------------------- // Public entry points @@ -38,6 +39,7 @@ pub fn all_controller_schemas() -> Vec { out.extend(kv_graph::FUNCTIONS.iter().map(|f| schemas(f))); out.extend(sync::FUNCTIONS.iter().map(|f| schemas(f))); out.extend(learn::FUNCTIONS.iter().map(|f| schemas(f))); + out.extend(tool_memory::FUNCTIONS.iter().map(|f| schemas(f))); out } @@ -49,6 +51,7 @@ pub fn all_registered_controllers() -> Vec { out.extend(kv_graph::controllers()); out.extend(sync::controllers()); out.extend(learn::controllers()); + out.extend(tool_memory::controllers()); out } @@ -69,6 +72,9 @@ pub fn schemas(function: &str) -> ControllerSchema { if let Some(schema) = learn::schema(function) { return schema; } + if let Some(schema) = tool_memory::schema(function) { + return schema; + } unknown_schema() } diff --git a/src/openhuman/memory/schemas/tool_memory.rs b/src/openhuman/memory/schemas/tool_memory.rs new file mode 100644 index 000000000..832e038cc --- /dev/null +++ b/src/openhuman/memory/schemas/tool_memory.rs @@ -0,0 +1,269 @@ +//! Schemas and handlers for the tool-scoped memory RPC methods. +//! +//! Six methods are exposed under the `memory` namespace: +//! +//! - `tool_rule_put` — upsert a rule +//! - `tool_rule_get` — fetch a single rule by id +//! - `tool_rule_list` — list every rule for a tool +//! - `tool_rule_delete` — delete a rule +//! - `tool_rules_for_prompt` — pre-rendered prompt block + structured rules +//! - `tool_rules_json` — raw JSON list for envelope consumers + +use serde_json::{Map, Value}; + +use crate::core::all::{ControllerFuture, RegisteredController}; +use crate::core::{ControllerSchema, FieldSchema, TypeSchema}; +use crate::openhuman::memory::rpc::{ + self, ToolRuleListParams, ToolRulePutParams, ToolRuleRefParams, ToolRulesForPromptParams, +}; + +use super::{parse_params, to_json}; + +pub(super) const FUNCTIONS: &[&str] = &[ + "tool_rule_put", + "tool_rule_get", + "tool_rule_list", + "tool_rule_delete", + "tool_rules_for_prompt", + "tool_rules_json", +]; + +pub(super) fn controllers() -> Vec { + vec![ + RegisteredController { + schema: schema("tool_rule_put").unwrap(), + handler: handle_tool_rule_put, + }, + RegisteredController { + schema: schema("tool_rule_get").unwrap(), + handler: handle_tool_rule_get, + }, + RegisteredController { + schema: schema("tool_rule_list").unwrap(), + handler: handle_tool_rule_list, + }, + RegisteredController { + schema: schema("tool_rule_delete").unwrap(), + handler: handle_tool_rule_delete, + }, + RegisteredController { + schema: schema("tool_rules_for_prompt").unwrap(), + handler: handle_tool_rules_for_prompt, + }, + RegisteredController { + schema: schema("tool_rules_json").unwrap(), + handler: handle_tool_rules_json, + }, + ] +} + +pub(super) fn schema(function: &str) -> Option { + Some(match function { + "tool_rule_put" => ControllerSchema { + namespace: "memory", + function: "tool_rule_put", + description: + "Upsert a tool-scoped memory rule. Stored in namespace `tool-{tool_name}`, separate \ + from generic `global` / `skill-{id}` / `tool_effectiveness` namespaces. Use \ + `priority='critical'` for safety-critical rules that must survive context compression.", + inputs: vec![ + FieldSchema { + name: "tool_name", + ty: TypeSchema::String, + comment: "Tool the rule applies to.", + required: true, + }, + FieldSchema { + name: "rule", + ty: TypeSchema::String, + comment: "Natural-language rule body.", + required: true, + }, + FieldSchema { + name: "priority", + ty: TypeSchema::Option(Box::new(TypeSchema::String)), + comment: "Priority — 'critical', 'high', or 'normal'. Defaults to 'normal'.", + required: false, + }, + FieldSchema { + name: "source", + ty: TypeSchema::Option(Box::new(TypeSchema::String)), + comment: + "Provenance — 'user_explicit', 'post_turn', or 'programmatic'. Defaults to \ + 'programmatic'.", + required: false, + }, + FieldSchema { + name: "tags", + ty: TypeSchema::Array(Box::new(TypeSchema::String)), + comment: "Optional tags for filtering.", + required: false, + }, + FieldSchema { + name: "id", + ty: TypeSchema::Option(Box::new(TypeSchema::String)), + comment: "Optional rule id — supplied to upsert in place.", + required: false, + }, + ], + outputs: vec![FieldSchema { + name: "result", + ty: TypeSchema::Json, + comment: "Stored rule, including assigned id and timestamps.", + required: true, + }], + }, + "tool_rule_get" => ControllerSchema { + namespace: "memory", + function: "tool_rule_get", + description: "Fetch a tool-scoped rule by `(tool_name, id)`. Returns null when missing.", + inputs: vec![ + FieldSchema { + name: "tool_name", + ty: TypeSchema::String, + comment: "Tool the rule applies to.", + required: true, + }, + FieldSchema { + name: "id", + ty: TypeSchema::String, + comment: "Rule id.", + required: true, + }, + ], + outputs: vec![FieldSchema { + name: "result", + ty: TypeSchema::Json, + comment: "Stored rule or null.", + required: true, + }], + }, + "tool_rule_list" => ControllerSchema { + namespace: "memory", + function: "tool_rule_list", + description: + "List every rule for a tool, sorted by priority (critical → high → normal) and then \ + by `updated_at` descending.", + inputs: vec![FieldSchema { + name: "tool_name", + ty: TypeSchema::String, + comment: "Tool to list rules for.", + required: true, + }], + outputs: vec![FieldSchema { + name: "result", + ty: TypeSchema::Json, + comment: "Array of stored rules.", + required: true, + }], + }, + "tool_rule_delete" => ControllerSchema { + namespace: "memory", + function: "tool_rule_delete", + description: + "Delete a tool-scoped rule. Returns true if the rule existed before deletion.", + inputs: vec![ + FieldSchema { + name: "tool_name", + ty: TypeSchema::String, + comment: "Tool the rule applies to.", + required: true, + }, + FieldSchema { + name: "id", + ty: TypeSchema::String, + comment: "Rule id.", + required: true, + }, + ], + outputs: vec![FieldSchema { + name: "result", + ty: TypeSchema::Bool, + comment: "True when the rule existed.", + required: true, + }], + }, + "tool_rules_for_prompt" => ControllerSchema { + namespace: "memory", + function: "tool_rules_for_prompt", + description: + "Pre-fetch Critical + High priority rules for prompt injection. Returns the \ + rendered Markdown block (ready for the system prompt) plus the structured rule \ + snapshot. Used by the session builder to surface rules during tool-selection and \ + pre-execution phases.", + inputs: vec![FieldSchema { + name: "tools", + ty: TypeSchema::Array(Box::new(TypeSchema::String)), + comment: "Constrain to these tools. Empty means scan every known tool namespace.", + required: false, + }], + outputs: vec![FieldSchema { + name: "result", + ty: TypeSchema::Json, + comment: "{ rendered: string, rules: ToolMemoryRule[] }.", + required: true, + }], + }, + "tool_rules_json" => ControllerSchema { + namespace: "memory", + function: "tool_rules_json", + description: + "Return the full rule list for a tool as raw JSON — useful for envelope consumers.", + inputs: vec![FieldSchema { + name: "tool_name", + ty: TypeSchema::String, + comment: "Tool to list rules for.", + required: true, + }], + outputs: vec![FieldSchema { + name: "result", + ty: TypeSchema::Json, + comment: "Array of stored rules as JSON.", + required: true, + }], + }, + _ => return None, + }) +} + +fn handle_tool_rule_put(params: Map) -> ControllerFuture { + Box::pin(async move { + let payload = parse_params::(params)?; + to_json(rpc::tool_rule_put(payload).await?) + }) +} + +fn handle_tool_rule_get(params: Map) -> ControllerFuture { + Box::pin(async move { + let payload = parse_params::(params)?; + to_json(rpc::tool_rule_get(payload).await?) + }) +} + +fn handle_tool_rule_list(params: Map) -> ControllerFuture { + Box::pin(async move { + let payload = parse_params::(params)?; + to_json(rpc::tool_rule_list(payload).await?) + }) +} + +fn handle_tool_rule_delete(params: Map) -> ControllerFuture { + Box::pin(async move { + let payload = parse_params::(params)?; + to_json(rpc::tool_rule_delete(payload).await?) + }) +} + +fn handle_tool_rules_for_prompt(params: Map) -> ControllerFuture { + Box::pin(async move { + let payload = parse_params::(params)?; + to_json(rpc::tool_rules_for_prompt(payload).await?) + }) +} + +fn handle_tool_rules_json(params: Map) -> ControllerFuture { + Box::pin(async move { + let payload = parse_params::(params)?; + to_json(rpc::tool_rules_json(payload).await?) + }) +} diff --git a/src/openhuman/memory/schemas_tests.rs b/src/openhuman/memory/schemas_tests.rs index d7e410768..9781c2eb2 100644 --- a/src/openhuman/memory/schemas_tests.rs +++ b/src/openhuman/memory/schemas_tests.rs @@ -33,6 +33,13 @@ const ALL_FUNCTIONS: &[&str] = &[ "sync_all", "learn_all", "ingestion_status", + // Tool-scoped memory (#1400) + "tool_rule_put", + "tool_rule_get", + "tool_rule_list", + "tool_rule_delete", + "tool_rules_for_prompt", + "tool_rules_json", ]; #[test] diff --git a/src/openhuman/memory/store/client.rs b/src/openhuman/memory/store/client.rs index b6d9b6f15..edec3ddb3 100644 --- a/src/openhuman/memory/store/client.rs +++ b/src/openhuman/memory/store/client.rs @@ -58,6 +58,15 @@ impl MemoryClient { std::sync::Arc::clone(&self.inner.conn) } + /// Returns an `Arc` handle backed by the same + /// [`UnifiedMemory`] this client wraps. Used by sub-systems that + /// want to build on top of the `Memory` trait (e.g. the + /// tool-scoped memory layer) without depending on the concrete + /// `MemoryClient` type or holding a reference to it. + pub fn memory_handle(&self) -> Arc { + Arc::clone(&self.inner) as Arc + } + /// Create a new local memory client using the default `.openhuman` directory. /// /// # Errors diff --git a/src/openhuman/memory/tool_memory/capture.rs b/src/openhuman/memory/tool_memory/capture.rs new file mode 100644 index 000000000..dff774376 --- /dev/null +++ b/src/openhuman/memory/tool_memory/capture.rs @@ -0,0 +1,481 @@ +//! Post-turn capture hook for tool-scoped memory. +//! +//! This hook complements the statistics-only [`ToolTrackerHook`] — +//! `tool_effectiveness` records *what happened* (counts, error patterns), +//! while [`ToolMemoryCaptureHook`] records *what to do about it* as +//! actionable [`ToolMemoryRule`]s in the tool-scoped namespace. +//! +//! Two capture paths fire automatically after every turn: +//! +//! 1. **User edicts** — phrases like `never `, +//! `don't …`, or `stop ing …` in the user message are +//! promoted to a `Critical` rule attached to the matching tool when +//! one of the turn's tool calls plausibly applies. This covers the +//! "never email Sarah" safety case from the spec. +//! +//! 2. **Repeated tool failures** — when a tool fails twice or more +//! within a single turn, a `Normal`-priority observation is captured +//! so the agent has a record next time it considers that tool. +//! +//! Both paths are conservative — they only fire on clear signals, and +//! the captured rule body always points back to the user's own words so +//! a reviewer can see exactly what triggered it. +//! +//! Captured rules are stored via [`ToolMemoryStore`] in the +//! `tool-{tool_name}` namespace, never in `global` or +//! `tool_effectiveness`. +//! +//! [`ToolTrackerHook`]: crate::openhuman::learning::ToolTrackerHook +//! [`ToolMemoryStore`]: super::store::ToolMemoryStore + +use std::collections::HashMap; +use std::sync::Arc; + +use async_trait::async_trait; + +use super::store::ToolMemoryStore; +use super::types::{ToolMemoryPriority, ToolMemorySource}; +use crate::openhuman::agent::hooks::{PostTurnHook, ToolCallRecord, TurnContext}; +use crate::openhuman::memory::Memory; + +/// Maximum length (chars) of the captured rule body — keeps malformed or +/// runaway input from bloating the namespace. +const MAX_RULE_LEN: usize = 240; + +/// Post-turn hook that captures durable tool-scoped rules. +pub struct ToolMemoryCaptureHook { + store: ToolMemoryStore, + enabled: bool, +} + +impl ToolMemoryCaptureHook { + /// Build a new capture hook backed by the given memory. + pub fn new(memory: Arc, enabled: bool) -> Self { + Self { + store: ToolMemoryStore::new(memory), + enabled, + } + } + + /// Build a hook directly over a [`ToolMemoryStore`] — useful for + /// tests and call sites that already hold a store. + pub fn from_store(store: ToolMemoryStore, enabled: bool) -> Self { + Self { store, enabled } + } + + /// Look at the user message and return any `Critical`-priority rule + /// patterns it contains, paired with the tool name they apply to. + /// + /// Pure / synchronous so it can be unit-tested without a memory + /// backend. + pub fn extract_user_edicts( + user_message: &str, + tool_calls: &[ToolCallRecord], + ) -> Vec<(String, String)> { + let trimmed = user_message.trim(); + if trimmed.is_empty() { + return Vec::new(); + } + let lower = trimmed.to_lowercase(); + // Only treat "stop" as an imperative edict when it appears at a + // sentence boundary (start of message or after ". "/"\n"), so routine + // phrases like "I want to stop working" don't trigger false captures. + let stop_imperative = + lower.starts_with("stop ") || lower.contains(". stop ") || lower.contains("\nstop "); + if !(lower.contains("never ") || lower.contains("don't ") || lower.contains("do not ")) + && !stop_imperative + { + return Vec::new(); + } + + // Default tool: the first tool that ran in the turn. When there + // were no tool calls we still want to capture user edicts so + // they survive into the next turn — those land under the + // `__unscoped__` tool name and the agent can refile them. + let default_tool = tool_calls + .first() + .map(|tc| tc.name.clone()) + .unwrap_or_else(|| "__unscoped__".to_string()); + + let mut out = Vec::new(); + for raw_line in trimmed.split(|c: char| matches!(c, '.' | '\n' | ';')) { + let line = raw_line.trim(); + if line.is_empty() { + continue; + } + let lower_line = line.to_lowercase(); + let is_edict = lower_line.starts_with("never ") + || lower_line.starts_with("don't ") + || lower_line.starts_with("do not ") + || lower_line.starts_with("stop ") + || lower_line.contains(" never ") + || lower_line.contains(" don't ") + || lower_line.contains(" do not "); + if !is_edict { + continue; + } + let body: String = line.chars().take(MAX_RULE_LEN).collect(); + if body.is_empty() { + continue; + } + let tool = + pick_tool_for_edict(&body, tool_calls).unwrap_or_else(|| default_tool.clone()); + out.push((tool, body)); + } + out + } + + /// Look at the tool-call records and return any (tool_name, body) + /// pairs that describe repeated failures worth pinning as a + /// `Normal`-priority observation. + /// + /// A tool counts when it failed two or more times in the turn — + /// transient one-off failures are ignored to keep the namespace + /// from filling with noise. + pub fn extract_repeated_failures(tool_calls: &[ToolCallRecord]) -> Vec<(String, String)> { + let mut tallies: HashMap<&str, (usize, Option<&str>)> = HashMap::new(); + for tc in tool_calls { + if tc.success { + continue; + } + let entry = tallies.entry(tc.name.as_str()).or_insert((0, None)); + entry.0 += 1; + if entry.1.is_none() { + entry.1 = Some(tc.output_summary.as_str()); + } + } + + let mut out = Vec::new(); + for (tool, (count, sample)) in tallies { + if count < 2 { + continue; + } + let body = match sample { + Some(sample) => format!( + "Tool failed {count} times in one turn ({sample}). Consider an alternative \ + approach before retrying." + ), + None => format!( + "Tool failed {count} times in one turn. Consider an alternative approach \ + before retrying." + ), + }; + out.push((tool.to_string(), body.chars().take(MAX_RULE_LEN).collect())); + } + out + } +} + +#[async_trait] +impl PostTurnHook for ToolMemoryCaptureHook { + fn name(&self) -> &str { + "tool_memory_capture" + } + + async fn on_turn_complete(&self, ctx: &TurnContext) -> anyhow::Result<()> { + if !self.enabled { + return Ok(()); + } + + for (tool, body) in Self::extract_user_edicts(&ctx.user_message, &ctx.tool_calls) { + log::debug!( + "[tool-memory] capturing user edict tool={tool} body_len={}", + body.len() + ); + if let Err(err) = self + .store + .record( + &tool, + &body, + ToolMemoryPriority::Critical, + ToolMemorySource::UserExplicit, + vec!["user-edict".into()], + ) + .await + { + log::warn!("[tool-memory] failed to capture user edict for {tool}: {err}"); + } + } + + for (tool, body) in Self::extract_repeated_failures(&ctx.tool_calls) { + log::debug!( + "[tool-memory] capturing repeated failure tool={tool} body_len={}", + body.len() + ); + if let Err(err) = self + .store + .record( + &tool, + &body, + ToolMemoryPriority::Normal, + ToolMemorySource::PostTurn, + vec!["repeated-failure".into()], + ) + .await + { + log::warn!( + "[tool-memory] failed to capture repeated-failure observation for {tool}: {err}" + ); + } + } + + Ok(()) + } +} + +/// Helper: emit a [`ToolMemoryRule`] preview without flooding logs with +/// raw user prose. +fn truncate_for_log(body: &str) -> String { + let mut out: String = body.chars().take(80).collect(); + if body.chars().count() > 80 { + out.push('…'); + } + out +} + +/// Best-effort match between a user edict and a tool that ran in the +/// turn. We look for the tool name appearing as a word in the edict; +/// when several match, the first call's tool wins. +fn pick_tool_for_edict(body: &str, tool_calls: &[ToolCallRecord]) -> Option { + if tool_calls.is_empty() { + return None; + } + let lower = body.to_lowercase(); + for tc in tool_calls { + let needle = tc.name.to_lowercase(); + if needle.is_empty() { + continue; + } + if lower.contains(&needle) { + return Some(tc.name.clone()); + } + // Common-noun aliases — match "email" to a tool named + // "send_email", "gmail_send", etc. + for alias in tool_aliases(&tc.name) { + if lower.contains(alias) { + return Some(tc.name.clone()); + } + } + } + None +} + +/// Map a tool name to a small set of common-noun aliases users would +/// say in plain English ("email", "shell", "browser", …). Kept tiny on +/// purpose — anything more ambitious belongs in an LLM extractor. +fn tool_aliases(tool_name: &str) -> Vec<&'static str> { + let lower = tool_name.to_lowercase(); + let mut out = Vec::new(); + if lower.contains("mail") { + out.push("email"); + out.push("mail"); + } + if lower.contains("shell") || lower.contains("bash") || lower.contains("exec") { + out.push("shell"); + out.push("terminal"); + } + if lower.contains("browser") || lower.contains("web") || lower.contains("http") { + out.push("browser"); + out.push("web"); + } + if lower.contains("slack") { + out.push("slack"); + out.push("dm"); + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::openhuman::agent::hooks::ToolCallRecord; + use crate::openhuman::memory::tool_memory::store::ToolMemoryStore; + use crate::openhuman::memory::tool_memory::test_helpers::MockMemory; + + fn ctx_with(message: &str, tool_calls: Vec) -> TurnContext { + TurnContext { + user_message: message.into(), + assistant_response: "ok".into(), + tool_calls, + turn_duration_ms: 1, + session_id: None, + iteration_count: 1, + } + } + + fn call(name: &str, success: bool) -> ToolCallRecord { + ToolCallRecord { + name: name.into(), + arguments: serde_json::json!({}), + success, + output_summary: if success { + "ok".into() + } else { + "permission denied".into() + }, + duration_ms: 10, + } + } + + #[test] + fn extract_user_edicts_picks_up_never_phrase() { + let edicts = ToolMemoryCaptureHook::extract_user_edicts( + "Never email Sarah at sarah@example.com — she does not want updates.", + &[call("send_email", true)], + ); + assert!(!edicts.is_empty(), "expected at least one captured edict"); + let (tool, body) = &edicts[0]; + assert_eq!( + tool, "send_email", + "should map 'email' alias to send_email tool" + ); + assert!(body.to_lowercase().contains("never email")); + } + + #[test] + fn extract_user_edicts_handles_dont_and_stop_phrases() { + let edicts = ToolMemoryCaptureHook::extract_user_edicts( + "Don't run shell commands with sudo. Stop using browser for that.", + &[call("shell", true), call("browser", true)], + ); + assert_eq!(edicts.len(), 2, "should capture each imperative separately"); + } + + #[test] + fn extract_user_edicts_returns_empty_when_no_edict_present() { + let edicts = ToolMemoryCaptureHook::extract_user_edicts( + "Send Sarah an update when you can.", + &[call("send_email", true)], + ); + assert!(edicts.is_empty()); + } + + #[test] + fn extract_user_edicts_falls_back_to_first_tool_when_no_alias_match() { + let edicts = ToolMemoryCaptureHook::extract_user_edicts( + "Never do that automatically.", + &[call("calendar", true)], + ); + assert_eq!(edicts.len(), 1); + assert_eq!(edicts[0].0, "calendar"); + } + + #[test] + fn extract_user_edicts_uses_sentinel_when_no_tools_ran() { + let edicts = ToolMemoryCaptureHook::extract_user_edicts("Never do that.", &[]); + assert_eq!(edicts.len(), 1); + assert_eq!(edicts[0].0, "__unscoped__"); + } + + #[test] + fn extract_repeated_failures_needs_two_or_more_failures() { + let observations = ToolMemoryCaptureHook::extract_repeated_failures(&[ + call("shell", false), + call("shell", false), + call("shell", true), + ]); + assert_eq!(observations.len(), 1); + assert_eq!(observations[0].0, "shell"); + assert!(observations[0].1.contains("failed 2 times")); + } + + #[test] + fn extract_repeated_failures_ignores_single_failures() { + let observations = + ToolMemoryCaptureHook::extract_repeated_failures(&[call("shell", false)]); + assert!(observations.is_empty()); + } + + #[tokio::test] + async fn on_turn_complete_persists_critical_rule_for_user_edict() { + let memory: Arc = Arc::new(MockMemory::default()); + let store = ToolMemoryStore::new(memory.clone()); + let hook = ToolMemoryCaptureHook::from_store(store.clone(), true); + + hook.on_turn_complete(&ctx_with( + "Never email Sarah — she opted out.", + vec![call("send_email", true)], + )) + .await + .unwrap(); + + let rules = store.list_rules("send_email").await.unwrap(); + assert_eq!(rules.len(), 1); + assert_eq!(rules[0].priority, ToolMemoryPriority::Critical); + assert_eq!(rules[0].source, ToolMemorySource::UserExplicit); + assert!(rules[0].tags.contains(&"user-edict".to_string())); + } + + #[tokio::test] + async fn on_turn_complete_no_op_when_disabled() { + let memory: Arc = Arc::new(MockMemory::default()); + let store = ToolMemoryStore::new(memory.clone()); + let hook = ToolMemoryCaptureHook::from_store(store.clone(), false); + hook.on_turn_complete(&ctx_with( + "Never email Sarah.", + vec![call("send_email", true)], + )) + .await + .unwrap(); + assert!(store.list_rules("send_email").await.unwrap().is_empty()); + } + + /// Safety case (AC #5): "never email Sarah" flows end-to-end from + /// a user utterance → captured as a Critical rule → surfaces in + /// the prompt-injection block. + #[tokio::test] + async fn safety_case_never_email_sarah_pins_into_prompt_block() { + let memory: Arc = Arc::new(MockMemory::default()); + let store = ToolMemoryStore::new(memory.clone()); + let hook = ToolMemoryCaptureHook::from_store(store.clone(), true); + + // 1. Capture the edict from a normal user turn. + hook.on_turn_complete(&ctx_with( + "Never email Sarah at sarah@example.com.", + vec![call("send_email", true)], + )) + .await + .unwrap(); + + // 2. The rule lands in the tool-scoped namespace with Critical + // priority — distinct from `tool_effectiveness` / global. + let stored = store.list_rules("send_email").await.unwrap(); + assert_eq!(stored.len(), 1); + assert_eq!(stored[0].priority, ToolMemoryPriority::Critical); + + // 3. `rules_for_prompt` pulls it eagerly so the session builder + // can pin it into the (compression-resistant) system prompt. + let prompt = store + .rules_for_prompt(&["send_email".to_string()]) + .await + .unwrap(); + assert!(prompt.contains_key("send_email")); + + // 4. The rendered block is non-empty and mentions the edict + // verbatim — the exact bytes the safety pipeline puts in + // front of the agent on every subsequent turn. + let mut flat: Vec<_> = prompt.into_values().flatten().collect(); + flat.sort_by(|a, b| b.priority.cmp(&a.priority)); + let rendered = crate::openhuman::memory::tool_memory::render_tool_memory_rules(&flat); + assert!(rendered.contains("Never email Sarah")); + assert!(rendered.contains("**[critical]**")); + } + + #[tokio::test] + async fn on_turn_complete_records_repeated_failure_observation() { + let memory: Arc = Arc::new(MockMemory::default()); + let store = ToolMemoryStore::new(memory.clone()); + let hook = ToolMemoryCaptureHook::from_store(store.clone(), true); + hook.on_turn_complete(&ctx_with( + "Try again", + vec![call("shell", false), call("shell", false)], + )) + .await + .unwrap(); + let rules = store.list_rules("shell").await.unwrap(); + assert_eq!(rules.len(), 1); + assert_eq!(rules[0].priority, ToolMemoryPriority::Normal); + assert_eq!(rules[0].source, ToolMemorySource::PostTurn); + assert!(rules[0].tags.contains(&"repeated-failure".to_string())); + } +} diff --git a/src/openhuman/memory/tool_memory/mod.rs b/src/openhuman/memory/tool_memory/mod.rs new file mode 100644 index 000000000..02845c3ed --- /dev/null +++ b/src/openhuman/memory/tool_memory/mod.rs @@ -0,0 +1,44 @@ +//! Tool-scoped memory layer for durable learnings and high-priority rules. +//! +//! Implements the dedicated memory namespace requested in +//! [issue #1400](https://github.com/tinyhumansai/openhuman/issues/1400): +//! a first-class storage and retrieval surface for **actionable** +//! tool-specific guidance, distinct from the +//! [`tool_effectiveness`](crate::openhuman::learning::tool_tracker) +//! statistics namespace and from the generic `global` / `skill-*` +//! namespaces. +//! +//! ## Namespace convention +//! +//! Each tool gets its own namespace `tool-{tool_name}`. The prefix is +//! distinct from `global`, `skill-{id}`, `tool_effectiveness`, and the +//! learning namespaces so list/clear operations can reason about it +//! without ambiguity. Build the namespace string via +//! [`types::tool_memory_namespace`] — never hard-code the format. +//! +//! ## Components +//! +//! - [`types`] — [`ToolMemoryRule`], [`ToolMemoryPriority`], +//! [`ToolMemorySource`]. +//! - [`store`] — [`ToolMemoryStore`], the put/list/delete/prompt API +//! built on top of an `Arc`. +//! - [`capture`] — [`ToolMemoryCaptureHook`], the post-turn +//! [`PostTurnHook`] that records user edicts and repeated tool +//! failures. +//! - [`prompt`] — [`ToolMemoryRulesSection`], the prompt section that +//! pins Critical / High rules into the system prompt so they survive +//! mid-session compression. +//! +//! [`PostTurnHook`]: crate::openhuman::agent::hooks::PostTurnHook + +pub mod capture; +pub mod prompt; +pub mod store; +#[cfg(test)] +pub mod test_helpers; +pub mod types; + +pub use capture::ToolMemoryCaptureHook; +pub use prompt::{render_tool_memory_rules, ToolMemoryRulesSection, TOOL_MEMORY_HEADING}; +pub use store::{ToolMemoryStore, TOOL_MEMORY_PROMPT_CAP}; +pub use types::{tool_memory_namespace, ToolMemoryPriority, ToolMemoryRule, ToolMemorySource}; diff --git a/src/openhuman/memory/tool_memory/prompt.rs b/src/openhuman/memory/tool_memory/prompt.rs new file mode 100644 index 000000000..3fc399bcb --- /dev/null +++ b/src/openhuman/memory/tool_memory/prompt.rs @@ -0,0 +1,250 @@ +//! Prompt section that injects tool-scoped memory rules into the +//! system prompt. +//! +//! ## Why a prompt section +//! +//! Mid-session compression rewrites the rolling chat buffer but never +//! the system prompt — that prompt is frozen for the whole session by +//! design (so the inference backend's prefix cache stays warm; see +//! [`crate::openhuman::agent::prompts::SystemPromptBuilder::build`]). +//! +//! Anything we want to be **compression-resistant** therefore has to +//! live in the system prompt. That is exactly where Critical and High +//! priority [`ToolMemoryRule`]s belong: a "never email Sarah" rule +//! cannot be silently dropped when the buffer fills up. +//! +//! ## What gets rendered +//! +//! The section takes ownership of the caller-supplied list of rules +//! (already filtered to the eager priorities by +//! [`ToolMemoryStore::rules_for_prompt`]) at construction time, mirrors +//! the pattern used by +//! [`crate::openhuman::agent::prompts::ReflectionMemoryContextSection`]. +//! Snapshot semantics — the rendered bytes are stable for the lifetime +//! of the session, preserving the inference backend's prefix cache hit. +//! +//! [`ToolMemoryRule`]: super::types::ToolMemoryRule +//! [`ToolMemoryStore::rules_for_prompt`]: super::store::ToolMemoryStore::rules_for_prompt + +use anyhow::Result; + +use super::types::{ToolMemoryPriority, ToolMemoryRule}; +use crate::openhuman::context::prompt::{PromptContext, PromptSection}; + +/// Heading injected when at least one rule is present. +pub const TOOL_MEMORY_HEADING: &str = "## Tool-scoped rules"; + +/// Prompt section that renders an at-construction snapshot of +/// [`ToolMemoryRule`]s into the system prompt. +/// +/// Construct via [`Self::new`] with the rules the session builder +/// pre-fetched from [`ToolMemoryStore::rules_for_prompt`]. +/// +/// [`ToolMemoryStore::rules_for_prompt`]: super::store::ToolMemoryStore::rules_for_prompt +pub struct ToolMemoryRulesSection { + rendered: String, +} + +impl ToolMemoryRulesSection { + /// Build a section from a pre-fetched rule snapshot. + /// + /// Rendering happens up-front so subsequent `build` calls — which + /// run once per system prompt assembly — are I/O-free and + /// deterministic. + pub fn new(rules: Vec) -> Self { + Self { + rendered: render_tool_memory_rules(&rules), + } + } + + /// Construct an empty section. Useful as a placeholder for builders + /// that always include the section name in their chain. + pub fn empty() -> Self { + Self { + rendered: String::new(), + } + } + + /// Returns true when the section will emit no output. + pub fn is_empty(&self) -> bool { + self.rendered.trim().is_empty() + } +} + +impl PromptSection for ToolMemoryRulesSection { + fn name(&self) -> &str { + "tool_memory_rules" + } + + fn build(&self, _ctx: &PromptContext<'_>) -> Result { + Ok(self.rendered.clone()) + } +} + +/// Pure rendering helper — public so callers that pre-render the block +/// (e.g. tests, dynamic prompt sources) can share the same logic. +pub fn render_tool_memory_rules(rules: &[ToolMemoryRule]) -> String { + if rules.is_empty() { + return String::new(); + } + + // Stable order: Critical first, then High; within a priority, by + // tool name, then by rule body. Callers may pass an already-sorted + // list (the store does), but rendering must not depend on that + // contract — the system prompt has to be byte-stable. + let mut sorted: Vec<&ToolMemoryRule> = rules.iter().collect(); + sorted.sort_by(|a, b| { + b.priority + .cmp(&a.priority) + .then_with(|| a.tool_name.cmp(&b.tool_name)) + .then_with(|| a.rule.cmp(&b.rule)) + .then_with(|| a.id.cmp(&b.id)) + }); + + let mut out = String::new(); + out.push_str(TOOL_MEMORY_HEADING); + out.push_str("\n\n"); + out.push_str( + "These rules are pinned by the user or by the safety pipeline. Treat \ + every entry as a hard constraint when considering the matching tool — \ + do not override them silently. Lower-priority guidance lives in the \ + `tool-{name}` memory namespace and can be queried via `memory_recall` \ + if needed.\n\n", + ); + + let mut current_tool: Option<&str> = None; + for rule in sorted { + if current_tool != Some(rule.tool_name.as_str()) { + if current_tool.is_some() { + out.push('\n'); + } + out.push_str("### `"); + out.push_str(rule.tool_name.as_str()); + out.push_str("`\n"); + current_tool = Some(rule.tool_name.as_str()); + } + out.push_str("- "); + out.push_str(priority_marker(rule.priority)); + out.push(' '); + out.push_str(rule.rule.trim()); + out.push('\n'); + } + + out +} + +fn priority_marker(priority: ToolMemoryPriority) -> &'static str { + match priority { + ToolMemoryPriority::Critical => "**[critical]**", + ToolMemoryPriority::High => "**[high]**", + ToolMemoryPriority::Normal => "**[normal]**", + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::openhuman::agent::prompts::types::{ + LearnedContextData, PromptContext, ToolCallFormat, + }; + use crate::openhuman::memory::tool_memory::types::ToolMemorySource; + + fn rule(tool: &str, body: &str, priority: ToolMemoryPriority) -> ToolMemoryRule { + ToolMemoryRule { + id: format!("{tool}/{body}"), + tool_name: tool.into(), + rule: body.into(), + priority, + source: ToolMemorySource::UserExplicit, + tags: vec![], + created_at: "2026-05-11T00:00:00Z".into(), + updated_at: "2026-05-11T00:00:00Z".into(), + } + } + + #[test] + fn renders_empty_when_no_rules() { + assert!(render_tool_memory_rules(&[]).is_empty()); + } + + #[test] + fn section_empty_returns_blank_build_output() { + let section = ToolMemoryRulesSection::empty(); + assert!(section.is_empty()); + } + + #[test] + fn renders_heading_and_priority_markers() { + let rules = vec![ + rule("email", "never email Sarah", ToolMemoryPriority::Critical), + rule("shell", "avoid sudo", ToolMemoryPriority::High), + ]; + let out = render_tool_memory_rules(&rules); + assert!(out.contains(TOOL_MEMORY_HEADING)); + assert!(out.contains("### `email`")); + assert!(out.contains("### `shell`")); + assert!(out.contains("**[critical]**")); + assert!(out.contains("**[high]**")); + assert!(out.contains("never email Sarah")); + assert!(out.contains("avoid sudo")); + } + + #[test] + fn renders_critical_before_high_regardless_of_input_order() { + let rules = vec![ + rule("shell", "avoid sudo", ToolMemoryPriority::High), + rule("email", "never email Sarah", ToolMemoryPriority::Critical), + ]; + let out = render_tool_memory_rules(&rules); + let critical_pos = out.find("never email Sarah").unwrap(); + let high_pos = out.find("avoid sudo").unwrap(); + assert!( + critical_pos < high_pos, + "Critical rules must render before High; output:\n{out}" + ); + } + + #[test] + fn renders_byte_stable_output_for_identical_inputs() { + let rules = vec![ + rule("email", "never email Sarah", ToolMemoryPriority::Critical), + rule("shell", "avoid sudo", ToolMemoryPriority::High), + ]; + let first = render_tool_memory_rules(&rules); + let again = render_tool_memory_rules(&rules); + assert_eq!(first, again); + } + + #[test] + fn section_renders_via_prompt_section_trait() { + // build() must not depend on PromptContext fields — it returns + // the at-construction snapshot verbatim. We call it here to + // exercise the trait contract directly. + let section = ToolMemoryRulesSection::new(vec![rule( + "email", + "never email Sarah", + ToolMemoryPriority::Critical, + )]); + assert!(!section.is_empty()); + let visible = std::collections::HashSet::new(); + let ctx = PromptContext { + workspace_dir: std::path::Path::new("."), + model_name: "test", + agent_id: "test", + tools: &[], + skills: &[], + dispatcher_instructions: "", + learned: LearnedContextData::default(), + visible_tool_names: &visible, + tool_call_format: ToolCallFormat::PFormat, + connected_integrations: &[], + connected_identities_md: String::new(), + include_profile: false, + include_memory_md: false, + curated_snapshot: None, + user_identity: None, + }; + let built = section.build(&ctx).unwrap(); + assert!(built.contains("never email Sarah")); + } +} diff --git a/src/openhuman/memory/tool_memory/store.rs b/src/openhuman/memory/tool_memory/store.rs new file mode 100644 index 000000000..db4fa8c0a --- /dev/null +++ b/src/openhuman/memory/tool_memory/store.rs @@ -0,0 +1,269 @@ +//! Storage layer for [`ToolMemoryRule`]s. +//! +//! Rules are persisted as KV rows in the tool's dedicated namespace +//! (`tool-{tool_name}`). KV storage is preferred over the document / +//! embedding pipeline because: +//! +//! - Tool guidance is short, structured, and benefits from exact key +//! lookup more than semantic retrieval. +//! - Writes never block on the local embedding model. +//! - Atomicity per rule is sufficient for safety-critical instructions. +//! +//! This module wraps an [`Arc`] handle (rather than the +//! concrete [`MemoryClient`]) so unit tests can swap in an in-memory +//! mock following the existing `tool_tracker` pattern. + +use std::collections::HashMap; +use std::sync::Arc; + +use serde_json::Value; + +use super::types::{tool_memory_namespace, ToolMemoryPriority, ToolMemoryRule, ToolMemorySource}; +use crate::openhuman::memory::{Memory, MemoryCategory}; + +/// Maximum number of rules surfaced into the system prompt at once. +/// +/// Keeps the cache-friendly prefix bounded even when callers stash a long +/// list of Critical rules over time. Lower-priority rules are still +/// available via [`ToolMemoryStore::list_rules`]. +pub const TOOL_MEMORY_PROMPT_CAP: usize = 30; + +/// High-level store for tool-scoped memory rules. +/// +/// All methods operate on a single shared [`Arc`] backend. +/// Cheap to clone — the backend is reference-counted. +#[derive(Clone)] +pub struct ToolMemoryStore { + memory: Arc, +} + +impl ToolMemoryStore { + /// Build a new store over the given memory backend. + pub fn new(memory: Arc) -> Self { + Self { memory } + } + + /// Upsert a rule and return the stored copy (with `updated_at` + /// refreshed). + /// + /// If a rule with the same `(tool_name, id)` already exists, its + /// `created_at` is preserved. `tool_name` is sourced from the rule + /// itself to avoid storage/namespace skew. + pub async fn put_rule(&self, mut rule: ToolMemoryRule) -> Result { + if rule.tool_name.trim().is_empty() { + return Err("tool_name is required".to_string()); + } + if rule.rule.trim().is_empty() { + return Err("rule body is required".to_string()); + } + if rule.id.trim().is_empty() { + rule.id = ToolMemoryRule::generate_id(); + } + + let namespace = tool_memory_namespace(&rule.tool_name); + let key = ToolMemoryRule::storage_key(&rule.id); + + // Preserve created_at on upsert. + if let Some(existing) = self.fetch_rule(&namespace, &key).await? { + rule.created_at = existing.created_at; + } + rule.updated_at = chrono::Utc::now().to_rfc3339(); + + let content = serde_json::to_string(&rule).map_err(|e| e.to_string())?; + self.memory + .store( + &namespace, + &key, + &content, + MemoryCategory::Custom("tool_memory".into()), + None, + ) + .await + .map_err(|e| format!("store tool rule: {e:#}"))?; + + log::debug!( + "[tool-memory] put rule id={} tool={} priority={:?} source={:?}", + rule.id, + rule.tool_name, + rule.priority, + rule.source + ); + + Ok(rule) + } + + /// Fetch a single rule by `(tool_name, id)`. + pub async fn get_rule( + &self, + tool_name: &str, + rule_id: &str, + ) -> Result, String> { + let namespace = tool_memory_namespace(tool_name); + let key = ToolMemoryRule::storage_key(rule_id); + self.fetch_rule(&namespace, &key).await + } + + /// List every rule registered for a tool, sorted by priority (high + /// first) and then `updated_at` descending. + pub async fn list_rules(&self, tool_name: &str) -> Result, String> { + let namespace = tool_memory_namespace(tool_name); + let entries = self + .memory + .list(Some(&namespace), None, None) + .await + .map_err(|e| format!("list tool rules: {e:#}"))?; + + let mut rules: Vec = entries + .into_iter() + .filter(|entry| entry.key.starts_with("rule/")) + .filter_map( + |entry| match serde_json::from_str::(&entry.content) { + Ok(rule) => Some(rule), + Err(err) => { + log::warn!( + "[tool-memory] skipping malformed rule key={} tool={tool_name}: {err}", + entry.key + ); + None + } + }, + ) + .collect(); + + rules.sort_by(|a, b| { + b.priority + .cmp(&a.priority) + .then_with(|| b.updated_at.cmp(&a.updated_at)) + }); + + Ok(rules) + } + + /// Delete a rule. Returns `true` if the rule existed. + pub async fn delete_rule(&self, tool_name: &str, rule_id: &str) -> Result { + let namespace = tool_memory_namespace(tool_name); + let key = ToolMemoryRule::storage_key(rule_id); + self.memory + .forget(&namespace, &key) + .await + .map_err(|e| format!("forget tool rule: {e:#}")) + } + + /// Returns the set of rules whose [`ToolMemoryPriority`] indicates + /// they must be eagerly surfaced (Critical + High), grouped by tool + /// name. Result is bounded by [`TOOL_MEMORY_PROMPT_CAP`] entries + /// total — Critical rules are always preferred over High when the + /// cap is reached. + /// + /// `tools` constrains which tool namespaces to inspect; passing an + /// empty slice scans every known tool namespace via + /// [`Memory::namespace_summaries`]. + pub async fn rules_for_prompt( + &self, + tools: &[String], + ) -> Result>, String> { + let tool_names = if tools.is_empty() { + self.list_tool_names().await? + } else { + tools + .iter() + .map(|name| name.trim().to_string()) + .filter(|name| !name.is_empty()) + .collect() + }; + + let mut collected: Vec = Vec::new(); + for tool in &tool_names { + let rules = self.list_rules(tool).await?; + collected.extend(rules.into_iter().filter(|r| r.priority.is_eager())); + } + + // Critical first, then High; within a priority, freshest first. + collected.sort_by(|a, b| { + b.priority + .cmp(&a.priority) + .then_with(|| b.updated_at.cmp(&a.updated_at)) + }); + collected.truncate(TOOL_MEMORY_PROMPT_CAP); + + let mut out: HashMap> = HashMap::new(); + for rule in collected { + out.entry(rule.tool_name.clone()).or_default().push(rule); + } + Ok(out) + } + + /// Enumerate every tool that has at least one stored rule, by + /// inspecting namespace summaries and keeping only the `tool-…` + /// prefixed ones. + pub async fn list_tool_names(&self) -> Result, String> { + let summaries = self + .memory + .namespace_summaries() + .await + .map_err(|e| format!("list tool namespaces: {e:#}"))?; + let mut out = Vec::new(); + for summary in summaries { + if let Some(tool) = summary.namespace.strip_prefix("tool-") { + // Exclude empty names and the sentinel used for unscoped + // edicts captured before any tool call ran — those rules are + // not permanently associated with a real tool and must not be + // injected into prompt filtering for arbitrary sessions. + if !tool.is_empty() && tool != "__unscoped__" { + out.push(tool.to_string()); + } + } + } + out.sort(); + out.dedup(); + Ok(out) + } + + /// Convenience constructor: build a rule from caller-supplied fields + /// and persist it. Returns the stored rule. + pub async fn record( + &self, + tool_name: &str, + rule_body: &str, + priority: ToolMemoryPriority, + source: ToolMemorySource, + tags: Vec, + ) -> Result { + let mut rule = ToolMemoryRule::new(tool_name, rule_body, priority, source); + rule.tags = tags; + self.put_rule(rule).await + } + + /// Render rules for a single tool into a JSON value suitable for + /// passing through the RPC envelope. Sorted by priority desc. + pub async fn list_rules_json(&self, tool_name: &str) -> Result { + let rules = self.list_rules(tool_name).await?; + serde_json::to_value(rules).map_err(|e| e.to_string()) + } + + async fn fetch_rule( + &self, + namespace: &str, + key: &str, + ) -> Result, String> { + let entry = self + .memory + .get(namespace, key) + .await + .map_err(|e| format!("get tool rule: {e:#}"))?; + match entry { + Some(entry) => match serde_json::from_str::(&entry.content) { + Ok(rule) => Ok(Some(rule)), + Err(err) => { + log::warn!("[tool-memory] malformed rule entry in {namespace}/{key}: {err}"); + Ok(None) + } + }, + None => Ok(None), + } + } +} + +#[cfg(test)] +#[path = "store_tests.rs"] +mod tests; diff --git a/src/openhuman/memory/tool_memory/store_tests.rs b/src/openhuman/memory/tool_memory/store_tests.rs new file mode 100644 index 000000000..8fc097923 --- /dev/null +++ b/src/openhuman/memory/tool_memory/store_tests.rs @@ -0,0 +1,375 @@ +//! Tests for [`ToolMemoryStore`] — exercise the put/list/delete/prompt +//! surface against an in-memory mock backend. + +use std::sync::Arc; + +use super::*; +use crate::openhuman::memory::tool_memory::test_helpers::MockMemory; +use crate::openhuman::memory::tool_memory::types::{ + ToolMemoryPriority, ToolMemoryRule, ToolMemorySource, +}; + +fn fresh_store() -> ToolMemoryStore { + ToolMemoryStore::new(Arc::new(MockMemory::default())) +} + +#[tokio::test] +async fn put_rule_rejects_blank_tool_name() { + let store = fresh_store(); + let rule = ToolMemoryRule::new( + " ", + "body", + ToolMemoryPriority::Normal, + ToolMemorySource::Programmatic, + ); + let err = store.put_rule(rule).await.unwrap_err(); + assert!(err.contains("tool_name")); +} + +#[tokio::test] +async fn put_rule_rejects_blank_body() { + let store = fresh_store(); + let rule = ToolMemoryRule::new( + "email", + " ", + ToolMemoryPriority::Normal, + ToolMemorySource::Programmatic, + ); + let err = store.put_rule(rule).await.unwrap_err(); + assert!(err.contains("body")); +} + +#[tokio::test] +async fn put_then_get_round_trip_returns_same_rule() { + let store = fresh_store(); + let rule = store + .record( + "email", + "never email Sarah", + ToolMemoryPriority::Critical, + ToolMemorySource::UserExplicit, + vec!["safety".into()], + ) + .await + .unwrap(); + let fetched = store.get_rule("email", &rule.id).await.unwrap().unwrap(); + assert_eq!(fetched.tool_name, "email"); + assert_eq!(fetched.rule, "never email Sarah"); + assert_eq!(fetched.priority, ToolMemoryPriority::Critical); + assert_eq!(fetched.source, ToolMemorySource::UserExplicit); + assert_eq!(fetched.tags, vec!["safety".to_string()]); +} + +#[tokio::test] +async fn put_rule_preserves_created_at_on_upsert() { + let store = fresh_store(); + let mut rule = store + .record( + "email", + "rule body", + ToolMemoryPriority::Normal, + ToolMemorySource::Programmatic, + vec![], + ) + .await + .unwrap(); + let created_at = rule.created_at.clone(); + // Mutate and re-put under the same id. + rule.rule = "updated rule body".into(); + // Sleep a tiny amount to give the timestamp string a chance to change. + tokio::time::sleep(std::time::Duration::from_millis(5)).await; + let updated = store.put_rule(rule.clone()).await.unwrap(); + assert_eq!(updated.created_at, created_at); + assert_ne!(updated.updated_at, created_at); + assert_eq!(updated.rule, "updated rule body"); +} + +#[tokio::test] +async fn list_rules_sorts_critical_first_then_freshest() { + let store = fresh_store(); + store + .record( + "email", + "older normal", + ToolMemoryPriority::Normal, + ToolMemorySource::Programmatic, + vec![], + ) + .await + .unwrap(); + // Tiny sleep to ensure different updated_at strings. + tokio::time::sleep(std::time::Duration::from_millis(5)).await; + let crit = store + .record( + "email", + "never email Sarah", + ToolMemoryPriority::Critical, + ToolMemorySource::UserExplicit, + vec![], + ) + .await + .unwrap(); + tokio::time::sleep(std::time::Duration::from_millis(5)).await; + let high = store + .record( + "email", + "double-check the recipient", + ToolMemoryPriority::High, + ToolMemorySource::PostTurn, + vec![], + ) + .await + .unwrap(); + let rules = store.list_rules("email").await.unwrap(); + assert_eq!(rules.len(), 3); + assert_eq!(rules[0].id, crit.id); + assert_eq!(rules[1].id, high.id); + assert_eq!(rules[2].rule, "older normal"); +} + +#[tokio::test] +async fn delete_rule_removes_only_target_rule() { + let store = fresh_store(); + let rule = store + .record( + "shell", + "never run sudo", + ToolMemoryPriority::Critical, + ToolMemorySource::UserExplicit, + vec![], + ) + .await + .unwrap(); + store + .record( + "shell", + "prefer tmux for long-running commands", + ToolMemoryPriority::Normal, + ToolMemorySource::Programmatic, + vec![], + ) + .await + .unwrap(); + let deleted = store.delete_rule("shell", &rule.id).await.unwrap(); + assert!(deleted); + let remaining = store.list_rules("shell").await.unwrap(); + assert_eq!(remaining.len(), 1); + assert_ne!(remaining[0].id, rule.id); +} + +#[tokio::test] +async fn delete_rule_returns_false_when_missing() { + let store = fresh_store(); + let deleted = store.delete_rule("shell", "does-not-exist").await.unwrap(); + assert!(!deleted); +} + +#[tokio::test] +async fn list_tool_names_returns_only_tool_prefixed_namespaces() { + // Direct-write a global memory entry so we can verify the filter. + let memory: Arc = Arc::new(MockMemory::default()); + memory + .store( + "global", + "noise", + "{}", + MemoryCategory::Custom("misc".into()), + None, + ) + .await + .unwrap(); + let store = ToolMemoryStore::new(memory); + store + .record( + "email", + "rule a", + ToolMemoryPriority::Critical, + ToolMemorySource::UserExplicit, + vec![], + ) + .await + .unwrap(); + store + .record( + "shell", + "rule b", + ToolMemoryPriority::High, + ToolMemorySource::PostTurn, + vec![], + ) + .await + .unwrap(); + let mut tools = store.list_tool_names().await.unwrap(); + tools.sort(); + assert_eq!(tools, vec!["email".to_string(), "shell".to_string()]); +} + +#[tokio::test] +async fn rules_for_prompt_returns_only_eager_priorities() { + let store = fresh_store(); + store + .record( + "email", + "never email Sarah", + ToolMemoryPriority::Critical, + ToolMemorySource::UserExplicit, + vec![], + ) + .await + .unwrap(); + store + .record( + "email", + "double-check recipient", + ToolMemoryPriority::High, + ToolMemorySource::PostTurn, + vec![], + ) + .await + .unwrap(); + store + .record( + "email", + "we use BCC for newsletters", + ToolMemoryPriority::Normal, + ToolMemorySource::Programmatic, + vec![], + ) + .await + .unwrap(); + + let rendered = store + .rules_for_prompt(&["email".to_string()]) + .await + .unwrap(); + let email_rules = rendered.get("email").expect("email rules present"); + assert_eq!(email_rules.len(), 2, "Normal rule must not be eager"); + assert!(email_rules + .iter() + .any(|r| r.priority == ToolMemoryPriority::Critical)); + assert!(email_rules + .iter() + .any(|r| r.priority == ToolMemoryPriority::High)); + assert!(email_rules + .iter() + .all(|r| r.priority != ToolMemoryPriority::Normal)); +} + +#[tokio::test] +async fn rules_for_prompt_scans_all_namespaces_when_caller_passes_empty_slice() { + let store = fresh_store(); + store + .record( + "email", + "never email Sarah", + ToolMemoryPriority::Critical, + ToolMemorySource::UserExplicit, + vec![], + ) + .await + .unwrap(); + store + .record( + "shell", + "never run sudo", + ToolMemoryPriority::Critical, + ToolMemorySource::UserExplicit, + vec![], + ) + .await + .unwrap(); + let rendered = store.rules_for_prompt(&[]).await.unwrap(); + assert!(rendered.contains_key("email")); + assert!(rendered.contains_key("shell")); +} + +#[tokio::test] +async fn rules_for_prompt_caps_results() { + let store = fresh_store(); + for idx in 0..(TOOL_MEMORY_PROMPT_CAP + 5) { + store + .record( + "email", + &format!("high rule {idx}"), + ToolMemoryPriority::High, + ToolMemorySource::Programmatic, + vec![], + ) + .await + .unwrap(); + } + let rendered = store + .rules_for_prompt(&["email".to_string()]) + .await + .unwrap(); + let count: usize = rendered.values().map(|v| v.len()).sum(); + assert_eq!(count, TOOL_MEMORY_PROMPT_CAP); +} + +#[tokio::test] +async fn list_rules_skips_malformed_entries() { + let memory: Arc = Arc::new(MockMemory::default()); + // Manually write a malformed entry under the tool namespace and a + // valid one alongside it, then confirm the bad entry is dropped. + memory + .store( + "tool-email", + "rule/bad", + "{not-valid-json", + MemoryCategory::Custom("tool_memory".into()), + None, + ) + .await + .unwrap(); + let store = ToolMemoryStore::new(memory); + store + .record( + "email", + "valid", + ToolMemoryPriority::High, + ToolMemorySource::Programmatic, + vec![], + ) + .await + .unwrap(); + let rules = store.list_rules("email").await.unwrap(); + assert_eq!(rules.len(), 1); + assert_eq!(rules[0].rule, "valid"); +} + +#[tokio::test] +async fn list_rules_json_serializes_payload_for_rpc_envelopes() { + let store = fresh_store(); + store + .record( + "email", + "rule body", + ToolMemoryPriority::High, + ToolMemorySource::Programmatic, + vec![], + ) + .await + .unwrap(); + let json = store.list_rules_json("email").await.unwrap(); + let arr = json.as_array().expect("expected an array"); + assert_eq!(arr.len(), 1); + assert!(arr[0] + .get("priority") + .and_then(|v| v.as_str()) + .map(|s| s == "high") + .unwrap_or(false)); +} + +#[tokio::test] +async fn put_rule_assigns_id_when_blank() { + let store = fresh_store(); + let mut rule = ToolMemoryRule::new( + "email", + "rule body", + ToolMemoryPriority::Normal, + ToolMemorySource::Programmatic, + ); + rule.id = String::new(); + let stored = store.put_rule(rule).await.unwrap(); + assert!(!stored.id.is_empty()); +} diff --git a/src/openhuman/memory/tool_memory/test_helpers.rs b/src/openhuman/memory/tool_memory/test_helpers.rs new file mode 100644 index 000000000..0b73b2fd5 --- /dev/null +++ b/src/openhuman/memory/tool_memory/test_helpers.rs @@ -0,0 +1,107 @@ +//! Shared test infrastructure for the tool-scoped memory layer. +//! +//! Only compiled under `#[cfg(test)]`. + +use std::collections::HashMap; + +use async_trait::async_trait; +use parking_lot::Mutex; + +use crate::openhuman::memory::{Memory, MemoryCategory, MemoryEntry, NamespaceSummary, RecallOpts}; + +/// Minimal in-memory [`Memory`] backend for unit tests. +/// +/// Stores entries in a `HashMap` keyed by `(namespace, key)`. All methods +/// that are not needed by the store/capture tests are no-ops. +#[derive(Default)] +pub struct MockMemory { + pub entries: Mutex>, +} + +#[async_trait] +impl Memory for MockMemory { + fn name(&self) -> &str { + "mock" + } + async fn store( + &self, + namespace: &str, + key: &str, + content: &str, + category: MemoryCategory, + session_id: Option<&str>, + ) -> anyhow::Result<()> { + self.entries.lock().insert( + (namespace.to_string(), key.to_string()), + MemoryEntry { + id: format!("{namespace}/{key}"), + key: key.to_string(), + content: content.to_string(), + namespace: Some(namespace.to_string()), + category, + timestamp: "now".into(), + session_id: session_id.map(str::to_string), + score: None, + }, + ); + Ok(()) + } + async fn recall( + &self, + _query: &str, + _limit: usize, + _opts: RecallOpts<'_>, + ) -> anyhow::Result> { + Ok(Vec::new()) + } + async fn get(&self, namespace: &str, key: &str) -> anyhow::Result> { + Ok(self + .entries + .lock() + .get(&(namespace.to_string(), key.to_string())) + .cloned()) + } + async fn list( + &self, + namespace: Option<&str>, + _category: Option<&MemoryCategory>, + _session_id: Option<&str>, + ) -> anyhow::Result> { + let lock = self.entries.lock(); + Ok(match namespace { + Some(ns) => lock + .iter() + .filter(|((n, _), _)| n == ns) + .map(|(_, v)| v.clone()) + .collect(), + None => lock.iter().map(|(_, v)| v.clone()).collect(), + }) + } + async fn forget(&self, namespace: &str, key: &str) -> anyhow::Result { + Ok(self + .entries + .lock() + .remove(&(namespace.to_string(), key.to_string())) + .is_some()) + } + async fn namespace_summaries(&self) -> anyhow::Result> { + let mut counts: HashMap = HashMap::new(); + for ((ns, _), _) in self.entries.lock().iter() { + *counts.entry(ns.clone()).or_default() += 1; + } + Ok(counts + .into_iter() + .map(|(namespace, count)| NamespaceSummary { + namespace, + count, + last_updated: None, + }) + .collect()) + } + async fn count(&self) -> anyhow::Result { + Ok(self.entries.lock().len()) + } + async fn health_check(&self) -> bool { + true + } +} diff --git a/src/openhuman/memory/tool_memory/types.rs b/src/openhuman/memory/tool_memory/types.rs new file mode 100644 index 000000000..4ee541355 --- /dev/null +++ b/src/openhuman/memory/tool_memory/types.rs @@ -0,0 +1,239 @@ +//! Domain types for the tool-scoped memory layer. +//! +//! A [`ToolMemoryRule`] is a durable, actionable instruction attached to a +//! specific tool (e.g. `email`, `shell`, `web_search`). Unlike the per-tool +//! statistics stored in the `tool_effectiveness` namespace, these rules +//! capture **guidance** — corrections, safety constraints, and learned +//! operational rules that the agent should obey when considering or +//! invoking that tool. +//! +//! Rules carry a [`ToolMemoryPriority`] level so the retrieval pipeline can +//! distinguish safety-critical instructions from soft suggestions: +//! +//! - [`ToolMemoryPriority::Critical`] — pinned into the system prompt and +//! therefore not subject to mid-session context compression. +//! - [`ToolMemoryPriority::High`] — surfaced alongside critical rules at +//! tool-selection time. +//! - [`ToolMemoryPriority::Normal`] — available on demand via the recall +//! APIs, but not eagerly injected. + +use serde::{Deserialize, Serialize}; + +/// Priority/criticality of a [`ToolMemoryRule`]. +/// +/// Used by both storage (to filter what is pinned into the system prompt) +/// and retrieval (to sort high-priority guidance ahead of advisory notes). +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ToolMemoryPriority { + /// Soft suggestion — surfaced on demand, not eagerly injected. + Normal, + /// Important guidance — eagerly injected at tool-selection time. + High, + /// Safety-critical rule — pinned into the (compression-resistant) + /// system prompt so it survives the agent's full session. + Critical, +} + +impl Default for ToolMemoryPriority { + fn default() -> Self { + Self::Normal + } +} + +impl ToolMemoryPriority { + /// True for priorities that must be eagerly surfaced to the agent + /// (Critical/High rules are both pinned into the system prompt and + /// prefetched at session start, so they survive context compression). + pub fn is_eager(self) -> bool { + matches!(self, Self::Critical | Self::High) + } +} + +/// Where a [`ToolMemoryRule`] originated from. +/// +/// Recorded for provenance and so consumers (UI / debugging) can tell user +/// edicts apart from auto-captured observations. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ToolMemorySource { + /// User explicitly asked the agent to remember this rule. + UserExplicit, + /// Captured automatically from a post-turn observation (tool failure, + /// repeated correction, etc.). + PostTurn, + /// Written by another subsystem (e.g. an integration provisioner). + Programmatic, +} + +impl Default for ToolMemorySource { + fn default() -> Self { + Self::Programmatic + } +} + +/// A single tool-scoped memory rule. +/// +/// Stored under the `tool-{tool_name}` namespace as a KV entry keyed by +/// `rule/{rule_id}`. The id is stable across updates so callers can +/// upsert by replaying the same id. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct ToolMemoryRule { + /// Stable identifier within `(tool_name)`. Generated by callers via + /// [`ToolMemoryRule::generate_id`] when one is not supplied. + pub id: String, + /// Tool this rule applies to (e.g. `email`, `shell`). + pub tool_name: String, + /// Natural-language guidance that should reach the agent. + pub rule: String, + /// Criticality level for retrieval and compression behaviour. + #[serde(default)] + pub priority: ToolMemoryPriority, + /// Where this rule came from. + #[serde(default)] + pub source: ToolMemorySource, + /// Optional free-form tags for filtering (e.g. `safety`, `permission`). + #[serde(default)] + pub tags: Vec, + /// RFC3339 timestamp of when the rule was first written. + pub created_at: String, + /// RFC3339 timestamp of the last update. + pub updated_at: String, +} + +impl ToolMemoryRule { + /// Build a new rule with a freshly generated id and `created_at` / + /// `updated_at` set to "now". + pub fn new( + tool_name: impl Into, + rule: impl Into, + priority: ToolMemoryPriority, + source: ToolMemorySource, + ) -> Self { + let now = chrono::Utc::now().to_rfc3339(); + Self { + id: Self::generate_id(), + tool_name: tool_name.into(), + rule: rule.into(), + priority, + source, + tags: Vec::new(), + created_at: now.clone(), + updated_at: now, + } + } + + /// Generate a fresh, opaque rule id. + pub fn generate_id() -> String { + uuid::Uuid::new_v4().to_string() + } + + /// Storage key used inside the tool namespace. + pub fn storage_key(id: &str) -> String { + format!("rule/{id}") + } +} + +/// Namespace string for a given tool. Trimmed and lower-cased so callers +/// can pass user-supplied tool names without leaking whitespace into +/// downstream queries. +/// +/// The `tool-` prefix is intentionally distinct from `global`, `skill-…` +/// and `tool_effectiveness` so retrieval and clearing operations can +/// reason about the namespace without ambiguity. +pub fn tool_memory_namespace(tool_name: &str) -> String { + format!("tool-{}", tool_name.trim().to_lowercase()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn priority_default_is_normal() { + assert_eq!(ToolMemoryPriority::default(), ToolMemoryPriority::Normal); + } + + #[test] + fn priority_ordering_puts_critical_above_high() { + assert!(ToolMemoryPriority::Critical > ToolMemoryPriority::High); + assert!(ToolMemoryPriority::High > ToolMemoryPriority::Normal); + } + + #[test] + fn priority_is_eager_for_high_and_critical_only() { + assert!(ToolMemoryPriority::Critical.is_eager()); + assert!(ToolMemoryPriority::High.is_eager()); + assert!(!ToolMemoryPriority::Normal.is_eager()); + } + + #[test] + fn priority_snake_case_serde() { + assert_eq!( + serde_json::to_string(&ToolMemoryPriority::Critical).unwrap(), + "\"critical\"" + ); + assert_eq!( + serde_json::to_string(&ToolMemoryPriority::Normal).unwrap(), + "\"normal\"" + ); + } + + #[test] + fn source_default_is_programmatic() { + assert_eq!(ToolMemorySource::default(), ToolMemorySource::Programmatic); + } + + #[test] + fn rule_new_fills_id_and_timestamps() { + let rule = ToolMemoryRule::new( + "email", + "never email Sarah", + ToolMemoryPriority::Critical, + ToolMemorySource::UserExplicit, + ); + assert!(!rule.id.is_empty()); + assert_eq!(rule.tool_name, "email"); + assert_eq!(rule.rule, "never email Sarah"); + assert_eq!(rule.priority, ToolMemoryPriority::Critical); + assert_eq!(rule.source, ToolMemorySource::UserExplicit); + assert!(rule.created_at == rule.updated_at); + } + + #[test] + fn rule_generate_id_produces_unique_values() { + let a = ToolMemoryRule::generate_id(); + let b = ToolMemoryRule::generate_id(); + assert_ne!(a, b); + } + + #[test] + fn rule_storage_key_uses_rule_prefix() { + assert_eq!(ToolMemoryRule::storage_key("abc"), "rule/abc"); + } + + #[test] + fn rule_serde_roundtrip_preserves_fields() { + let rule = ToolMemoryRule { + id: "id-1".into(), + tool_name: "shell".into(), + rule: "never run sudo".into(), + priority: ToolMemoryPriority::High, + source: ToolMemorySource::PostTurn, + tags: vec!["safety".into()], + created_at: "2026-05-11T00:00:00Z".into(), + updated_at: "2026-05-11T00:00:01Z".into(), + }; + let json = serde_json::to_string(&rule).unwrap(); + let back: ToolMemoryRule = serde_json::from_str(&json).unwrap(); + assert_eq!(back, rule); + } + + #[test] + fn namespace_uses_tool_prefix_and_trims_whitespace() { + assert_eq!(tool_memory_namespace("email"), "tool-email"); + assert_eq!(tool_memory_namespace(" shell "), "tool-shell"); + assert_eq!(tool_memory_namespace("Send_Email"), "tool-send_email"); + assert_eq!(tool_memory_namespace("WebSearch"), "tool-websearch"); + } +}