mirror of
https://github.com/tinyhumansai/openhuman.git
synced 2026-07-27 21:08:00 +00:00
feat(memory): add tool-scoped memory with durable rule capture, RPC APIs, and prompt pinning (#1487)
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
YellowSnnowmann
co-authored by
Claude Sonnet 4.6
Claude Sonnet 4.6
parent
ef998ca16c
commit
a2ced40ef9
@@ -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)
|
||||
|
||||
@@ -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 <verb> <noun>`, `don't <verb> ...`, `do not <verb> ...`, or `stop <verb>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.
|
||||
@@ -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)",
|
||||
|
||||
@@ -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<String> =
|
||||
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<dyn Memory>,
|
||||
tool_names: &[String],
|
||||
) -> Vec<crate::openhuman::memory::ToolMemoryRule> {
|
||||
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()
|
||||
}
|
||||
}
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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<crate::openhuman::memory::ToolMemoryRule>,
|
||||
) -> 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<dyn PromptSection> =
|
||||
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).
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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() {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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<ToolMemoryPriority>,
|
||||
/// Provenance — defaults to `programmatic` when omitted.
|
||||
#[serde(default)]
|
||||
pub source: Option<ToolMemorySource>,
|
||||
/// Optional tags for filtering.
|
||||
#[serde(default)]
|
||||
pub tags: Vec<String>,
|
||||
/// Optional rule id — when supplied, the call upserts in place
|
||||
/// rather than creating a new entry.
|
||||
#[serde(default)]
|
||||
pub id: Option<String>,
|
||||
}
|
||||
|
||||
/// 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<String>,
|
||||
}
|
||||
|
||||
async fn open_store() -> Result<ToolMemoryStore, String> {
|
||||
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<RpcOutcome<ToolMemoryRule>, 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<RpcOutcome<Option<ToolMemoryRule>>, 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<RpcOutcome<Vec<ToolMemoryRule>>, 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<RpcOutcome<bool>, 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<ToolMemoryRule>,
|
||||
}
|
||||
|
||||
/// Pre-fetch Critical + High priority rules for prompt injection.
|
||||
pub async fn tool_rules_for_prompt(
|
||||
params: ToolRulesForPromptParams,
|
||||
) -> Result<RpcOutcome<ToolRulesForPromptResult>, 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<ToolMemoryRule> = 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<RpcOutcome<Value>, 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"))
|
||||
}
|
||||
@@ -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<ControllerSchema> {
|
||||
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<RegisteredController> {
|
||||
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()
|
||||
}
|
||||
|
||||
|
||||
@@ -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<RegisteredController> {
|
||||
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<ControllerSchema> {
|
||||
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<String, Value>) -> ControllerFuture {
|
||||
Box::pin(async move {
|
||||
let payload = parse_params::<ToolRulePutParams>(params)?;
|
||||
to_json(rpc::tool_rule_put(payload).await?)
|
||||
})
|
||||
}
|
||||
|
||||
fn handle_tool_rule_get(params: Map<String, Value>) -> ControllerFuture {
|
||||
Box::pin(async move {
|
||||
let payload = parse_params::<ToolRuleRefParams>(params)?;
|
||||
to_json(rpc::tool_rule_get(payload).await?)
|
||||
})
|
||||
}
|
||||
|
||||
fn handle_tool_rule_list(params: Map<String, Value>) -> ControllerFuture {
|
||||
Box::pin(async move {
|
||||
let payload = parse_params::<ToolRuleListParams>(params)?;
|
||||
to_json(rpc::tool_rule_list(payload).await?)
|
||||
})
|
||||
}
|
||||
|
||||
fn handle_tool_rule_delete(params: Map<String, Value>) -> ControllerFuture {
|
||||
Box::pin(async move {
|
||||
let payload = parse_params::<ToolRuleRefParams>(params)?;
|
||||
to_json(rpc::tool_rule_delete(payload).await?)
|
||||
})
|
||||
}
|
||||
|
||||
fn handle_tool_rules_for_prompt(params: Map<String, Value>) -> ControllerFuture {
|
||||
Box::pin(async move {
|
||||
let payload = parse_params::<ToolRulesForPromptParams>(params)?;
|
||||
to_json(rpc::tool_rules_for_prompt(payload).await?)
|
||||
})
|
||||
}
|
||||
|
||||
fn handle_tool_rules_json(params: Map<String, Value>) -> ControllerFuture {
|
||||
Box::pin(async move {
|
||||
let payload = parse_params::<ToolRuleListParams>(params)?;
|
||||
to_json(rpc::tool_rules_json(payload).await?)
|
||||
})
|
||||
}
|
||||
@@ -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]
|
||||
|
||||
@@ -58,6 +58,15 @@ impl MemoryClient {
|
||||
std::sync::Arc::clone(&self.inner.conn)
|
||||
}
|
||||
|
||||
/// Returns an `Arc<dyn Memory>` 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<dyn crate::openhuman::memory::Memory> {
|
||||
Arc::clone(&self.inner) as Arc<dyn crate::openhuman::memory::Memory>
|
||||
}
|
||||
|
||||
/// Create a new local memory client using the default `.openhuman` directory.
|
||||
///
|
||||
/// # Errors
|
||||
|
||||
@@ -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 <verb> <object>`,
|
||||
//! `don't <verb> …`, or `stop <verb>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<dyn Memory>, 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<String> {
|
||||
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<ToolCallRecord>) -> 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<dyn Memory> = 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<dyn Memory> = 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<dyn Memory> = 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<dyn Memory> = 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()));
|
||||
}
|
||||
}
|
||||
@@ -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<dyn Memory>`.
|
||||
//! - [`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};
|
||||
@@ -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<ToolMemoryRule>) -> 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<String> {
|
||||
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"));
|
||||
}
|
||||
}
|
||||
@@ -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<dyn Memory>`] 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<dyn Memory>`] backend.
|
||||
/// Cheap to clone — the backend is reference-counted.
|
||||
#[derive(Clone)]
|
||||
pub struct ToolMemoryStore {
|
||||
memory: Arc<dyn Memory>,
|
||||
}
|
||||
|
||||
impl ToolMemoryStore {
|
||||
/// Build a new store over the given memory backend.
|
||||
pub fn new(memory: Arc<dyn Memory>) -> 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<ToolMemoryRule, String> {
|
||||
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<Option<ToolMemoryRule>, 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<Vec<ToolMemoryRule>, 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<ToolMemoryRule> = entries
|
||||
.into_iter()
|
||||
.filter(|entry| entry.key.starts_with("rule/"))
|
||||
.filter_map(
|
||||
|entry| match serde_json::from_str::<ToolMemoryRule>(&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<bool, String> {
|
||||
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<HashMap<String, Vec<ToolMemoryRule>>, 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<ToolMemoryRule> = 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<String, Vec<ToolMemoryRule>> = 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<Vec<String>, 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<String>,
|
||||
) -> Result<ToolMemoryRule, String> {
|
||||
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<Value, String> {
|
||||
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<Option<ToolMemoryRule>, 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::<ToolMemoryRule>(&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;
|
||||
@@ -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<dyn Memory> = 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<dyn Memory> = 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());
|
||||
}
|
||||
@@ -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<HashMap<(String, String), MemoryEntry>>,
|
||||
}
|
||||
|
||||
#[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<Vec<MemoryEntry>> {
|
||||
Ok(Vec::new())
|
||||
}
|
||||
async fn get(&self, namespace: &str, key: &str) -> anyhow::Result<Option<MemoryEntry>> {
|
||||
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<Vec<MemoryEntry>> {
|
||||
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<bool> {
|
||||
Ok(self
|
||||
.entries
|
||||
.lock()
|
||||
.remove(&(namespace.to_string(), key.to_string()))
|
||||
.is_some())
|
||||
}
|
||||
async fn namespace_summaries(&self) -> anyhow::Result<Vec<NamespaceSummary>> {
|
||||
let mut counts: HashMap<String, usize> = 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<usize> {
|
||||
Ok(self.entries.lock().len())
|
||||
}
|
||||
async fn health_check(&self) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
@@ -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<String>,
|
||||
/// 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<String>,
|
||||
rule: impl Into<String>,
|
||||
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");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user