Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

memoryサーチツールを実装

Browse Source
This commit is contained in:
Hare
2026-04-27 17:24:08 +09:00
parent c8fe95901c
commit c9d1d4fa5c
10 changed files with 936 additions and 208 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/memory/src/tool/edit.rs
+44 -89
View File
@@ -1,11 +1,11 @@
//! `MemoryEdit` tool — partial string replacement on an existing memory record.
//!
//! Reads current content, applies the replacement, runs the Linter on
//! the result, writes only on success. The current-then-write window
//! is single-tool-call narrow; an external tracker is intentionally
//! omitted (memory tools are self-contained, no `tools` crate dep).
//! Reads current content by `(kind, slug)`, applies the replacement,
//! runs the Linter on the result, writes only on success. The
//! current-then-write window is single-tool-call narrow; an external
//! tracker is intentionally omitted (memory tools are self-contained,
//! no `tools` crate dep).
use std::path::PathBuf;
use std::sync::Arc;
use async_trait::async_trait;
@@ -13,18 +13,21 @@ use llm_worker::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use serde::Deserialize;
use crate::linter::{LintReport, Linter, WriteMode};
use crate::tool::MemoryToolKind;
use crate::workspace::WorkspaceLayout;
const DESCRIPTION: &str = "Replace a substring in an existing memory or knowledge \
record file. By default `old_string` must be unique in the file; set \
`replace_all: true` to replace every occurrence. The resulting content is \
re-validated by the memory linter; failure leaves the file untouched. Path \
must be absolute and lie inside the workspace's `memory/` or `knowledge/` tree.";
record selected by `kind` + `slug`. By default `old_string` must be unique in the \
file; set `replace_all: true` to replace every occurrence. The resulting content \
is re-validated by the memory linter; failure leaves the file untouched.";
#[derive(Debug, Deserialize, schemars::JsonSchema)]
struct EditParams {
/// Absolute path under the workspace's `memory/` or `knowledge/` tree.
file_path: PathBuf,
/// Record kind: `summary` | `decision` | `request` | `knowledge`.
kind: MemoryToolKind,
/// Slug. Required for everything except `summary`; forbidden for `summary`.
#[serde(default)]
slug: Option<String>,
/// String to replace. Must be unique in the file unless `replace_all` is true.
old_string: String,
/// Replacement string. Must differ from `old_string`.
@@ -35,6 +38,7 @@ struct EditParams {
}
struct EditTool {
layout: WorkspaceLayout,
linter: Linter,
}
@@ -45,12 +49,6 @@ impl Tool for EditTool {
ToolError::InvalidArgument(format!("invalid MemoryEdit input: {e}"))
})?;
if !params.file_path.is_absolute() {
return Err(ToolError::InvalidArgument(format!(
"file_path must be absolute: {}",
params.file_path.display()
)));
}
if params.old_string.is_empty() {
return Err(ToolError::InvalidArgument(
"old_string must not be empty".into(),
@@ -62,49 +60,30 @@ impl Tool for EditTool {
));
}
// Path-shape check; the layout::classify also runs inside the
// linter but we want a crisp error before reading the file.
if self
.linter
.layout()
.classify(&params.file_path)
.map_err(|e| ToolError::InvalidArgument(e.to_string()))?
.is_none()
{
return Err(ToolError::InvalidArgument(format!(
"path is not under the memory tree: {}",
params.file_path.display()
)));
}
let path = params.kind.resolve_path(&self.layout, params.slug.as_deref())?;
let current_bytes = std::fs::read(&params.file_path).map_err(|e| match e.kind() {
let current_bytes = std::fs::read(&path).map_err(|e| match e.kind() {
std::io::ErrorKind::NotFound => ToolError::ExecutionFailed(format!(
"file not found (use MemoryWrite to create): {}",
params.file_path.display()
)),
_ => ToolError::ExecutionFailed(format!(
"read failed at {}: {e}",
params.file_path.display()
"record not found (use MemoryWrite to create): {}",
path.display()
)),
_ => ToolError::ExecutionFailed(format!("read failed at {}: {e}", path.display())),
})?;
let current_text = std::str::from_utf8(&current_bytes).map_err(|_| {
ToolError::InvalidArgument(format!(
"file is not valid UTF-8: {}",
params.file_path.display()
))
ToolError::InvalidArgument(format!("file is not valid UTF-8: {}", path.display()))
})?;
let count = current_text.matches(&params.old_string).count();
if count == 0 {
return Err(ToolError::InvalidArgument(format!(
"old_string not found in {}",
params.file_path.display()
path.display()
)));
}
if !params.replace_all && count > 1 {
return Err(ToolError::InvalidArgument(format!(
"old_string occurs {count} times in {}; pass replace_all: true or narrow the snippet",
params.file_path.display()
path.display()
)));
}
@@ -115,21 +94,18 @@ impl Tool for EditTool {
};
let occurrences = if params.replace_all { count } else { 1 };
let report = self.linter.lint(&params.file_path, &new_text, WriteMode::Update);
let report = self.linter.lint(&path, &new_text, WriteMode::Update);
if report.has_errors() {
return Err(ToolError::InvalidArgument(format_report(&report)));
}
std::fs::write(&params.file_path, new_text.as_bytes()).map_err(|e| {
ToolError::ExecutionFailed(format!(
"failed to write {}: {e}",
params.file_path.display()
))
std::fs::write(&path, new_text.as_bytes()).map_err(|e| {
ToolError::ExecutionFailed(format!("failed to write {}: {e}", path.display()))
})?;
let summary = format!(
"Edited {} ({} replacement{}){}",
params.file_path.display(),
path.display(),
occurrences,
if occurrences == 1 { "" } else { "s" },
warning_tail(&report),
@@ -176,6 +152,7 @@ pub fn edit_tool(layout: WorkspaceLayout) -> ToolDefinition {
.description(DESCRIPTION)
.input_schema(schema_value);
let tool: Arc<dyn Tool> = Arc::new(EditTool {
layout: layout.clone(),
linter: Linter::new(layout.clone()),
});
(meta, tool)
@@ -186,6 +163,7 @@ pub fn edit_tool(layout: WorkspaceLayout) -> ToolDefinition {
mod tests {
use super::*;
use chrono::Utc;
use std::path::PathBuf;
use tempfile::TempDir;
fn now() -> String {
@@ -212,7 +190,8 @@ mod tests {
assert_eq!(meta.name, "MemoryEdit");
let inp = serde_json::json!({
"file_path": path.to_str().unwrap(),
"kind": "decision",
"slug": "foo",
"old_string": "body body",
"new_string": "edited",
});
@@ -230,7 +209,8 @@ mod tests {
// Drop the `status` field by replacing it with nothing.
let inp = serde_json::json!({
"file_path": path.to_str().unwrap(),
"kind": "decision",
"slug": "foo",
"old_string": "status: open\n",
"new_string": "",
});
@@ -244,12 +224,12 @@ mod tests {
}
#[tokio::test]
async fn edit_missing_file() {
let (dir, layout, _) = setup();
let other = dir.path().join("memory/decisions/ghost.md");
async fn edit_missing_record() {
let (_dir, layout, _) = setup();
let (_, tool) = edit_tool(layout)();
let inp = serde_json::json!({
"file_path": other.to_str().unwrap(),
"kind": "decision",
"slug": "ghost",
"old_string": "x",
"new_string": "y",
});
@@ -258,42 +238,17 @@ mod tests {
}
#[tokio::test]
async fn edit_outside_memory_tree_rejected() {
let (dir, layout, _) = setup();
let other = dir.path().join("src/lib.rs");
std::fs::create_dir_all(other.parent().unwrap()).unwrap();
std::fs::write(&other, "fn main() {}").unwrap();
async fn edit_workflow_kind_rejected() {
// Workflow is not exposed via MemoryToolKind, so deserialization fails.
let (_dir, layout, _) = setup();
let (_, tool) = edit_tool(layout)();
let inp = serde_json::json!({
"file_path": other.to_str().unwrap(),
"old_string": "fn",
"new_string": "pub fn",
"kind": "workflow",
"slug": "wf",
"old_string": "x",
"new_string": "y",
});
let err = tool.execute(&inp.to_string()).await.unwrap_err();
assert!(matches!(err, ToolError::InvalidArgument(_)));
}
#[tokio::test]
async fn edit_workflow_path_rejected() {
let (dir, layout, _) = setup();
let path = dir.path().join("memory/workflow/wf.md");
std::fs::create_dir_all(path.parent().unwrap()).unwrap();
let initial = format!(
"---\nupdated_at: {n}\ndescription: x\nauto_invoke: false\nuser_invocable: true\n---\nbody\n",
n = now()
);
std::fs::write(&path, &initial).unwrap();
let (_, tool) = edit_tool(layout)();
let inp = serde_json::json!({
"file_path": path.to_str().unwrap(),
"old_string": "body",
"new_string": "edited",
});
let err = tool.execute(&inp.to_string()).await.unwrap_err();
let msg = format!("{err}");
assert!(msg.to_lowercase().contains("workflow"), "{msg}");
// Original untouched.
assert!(std::fs::read_to_string(&path).unwrap().contains("body"));
}
}
crates/memory/src/tool/mod.rs
+82 -1
View File
@@ -1,9 +1,90 @@
//! Tool implementations stub. Filled in once the linter compiles green.
//! Memory-scoped tools: Read / Write / Edit / Search.
//!
//! All four take `kind` + `slug` (Summary takes only `kind`) and
//! resolve the path through [`WorkspaceLayout`]. The agent never has
//! to know the on-disk layout — Search returns `{slug, kind, ...}` and
//! that pair feeds straight into Read / Edit.
mod edit;
mod read;
mod search;
mod write;
use std::path::PathBuf;
use llm_worker::tool::ToolError;
use serde::Deserialize;
use crate::slug::Slug;
use crate::workspace::{RecordKind, WorkspaceLayout};
pub use edit::edit_tool;
pub use read::read_tool;
pub use search::{knowledge_search_tool, memory_search_tool, SearchConfig};
pub use write::write_tool;
/// Kinds the memory tools accept as input. `Workflow` is intentionally
/// excluded — workflows are sub-Worker context, not agent-editable.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, schemars::JsonSchema)]
#[serde(rename_all = "lowercase")]
pub enum MemoryToolKind {
Summary,
Decision,
Request,
Knowledge,
}
impl MemoryToolKind {
pub fn as_str(self) -> &'static str {
match self {
Self::Summary => "summary",
Self::Decision => "decision",
Self::Request => "request",
Self::Knowledge => "knowledge",
}
}
pub fn record_kind(self) -> RecordKind {
match self {
Self::Summary => RecordKind::Summary,
Self::Decision => RecordKind::Decision,
Self::Request => RecordKind::Request,
Self::Knowledge => RecordKind::Knowledge,
}
}
/// Resolve `(kind, slug)` to an absolute path under the workspace.
/// Summary forbids a slug; the per-record kinds require one.
pub fn resolve_path(
self,
layout: &WorkspaceLayout,
slug: Option<&str>,
) -> Result<PathBuf, ToolError> {
match self {
Self::Summary => {
if slug.is_some() {
return Err(ToolError::InvalidArgument(
"kind=summary does not accept a slug".to_string(),
));
}
Ok(layout.summary_path())
}
other => {
let raw = slug.ok_or_else(|| {
ToolError::InvalidArgument(format!(
"kind={} requires `slug`",
other.as_str()
))
})?;
let parsed = Slug::parse(raw)
.map_err(|e| ToolError::InvalidArgument(e.to_string()))?;
Ok(match other {
Self::Decision => layout.decision_path(&parsed),
Self::Request => layout.request_path(&parsed),
Self::Knowledge => layout.knowledge_path(&parsed),
Self::Summary => unreachable!(),
})
}
}
}
}
crates/memory/src/tool/read.rs
+71 -47
View File
@@ -1,29 +1,33 @@
//! `MemoryRead` tool.
//!
//! Constrained to `<workspace>/memory/` and `<workspace>/knowledge/`
//! paths. Returns line-numbered content (1-based), like the generic
//! Read tool, but rejects anything outside the memory tree so the
//! agent can't sneak in a non-memory read through this surface.
//! Reads a memory or knowledge record by `(kind, slug)`. Returns
//! line-numbered content (1-based), like the generic Read tool. The
//! agent never names a path — `Search` returns `{kind, slug, ...}`
//! and that pair feeds straight into Read.
use std::path::PathBuf;
use std::sync::Arc;
use async_trait::async_trait;
use llm_worker::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use serde::Deserialize;
use crate::tool::MemoryToolKind;
use crate::workspace::WorkspaceLayout;
const DESCRIPTION: &str = "Read a memory or knowledge record file under the \
workspace's `memory/` or `knowledge/` tree. Returns line-numbered output \
(1-based). Paths must be absolute and lie inside the memory tree.";
const DESCRIPTION: &str = "Read a memory or knowledge record by `kind` + `slug`. \
`kind` is one of: summary, decision, request, knowledge. \
For `summary` omit `slug`; for the others `slug` is required. \
Returns line-numbered output (1-based).";
const DEFAULT_LIMIT: usize = 2000;
#[derive(Debug, Deserialize, schemars::JsonSchema)]
struct ReadParams {
/// Absolute path to a file under the workspace's `memory/` or `knowledge/` tree.
file_path: PathBuf,
/// Record kind: `summary` | `decision` | `request` | `knowledge`.
kind: MemoryToolKind,
/// Slug. Required for everything except `summary`; forbidden for `summary`.
#[serde(default)]
slug: Option<String>,
/// 0-based line offset from the start. Defaults to 0.
#[serde(default)]
offset: Option<usize>,
@@ -43,33 +47,13 @@ impl Tool for ReadTool {
ToolError::InvalidArgument(format!("invalid MemoryRead input: {e}"))
})?;
if !params.file_path.is_absolute() {
return Err(ToolError::InvalidArgument(format!(
"file_path must be absolute: {}",
params.file_path.display()
)));
}
if self
.layout
.classify(&params.file_path)
.map_err(|e| ToolError::InvalidArgument(e.to_string()))?
.is_none()
{
return Err(ToolError::InvalidArgument(format!(
"path is not under the memory tree: {}",
params.file_path.display()
)));
}
let path = params.kind.resolve_path(&self.layout, params.slug.as_deref())?;
let bytes = std::fs::read(&params.file_path).map_err(|e| match e.kind() {
std::io::ErrorKind::NotFound => ToolError::ExecutionFailed(format!(
"file not found: {}",
params.file_path.display()
)),
_ => ToolError::ExecutionFailed(format!(
"read failed at {}: {e}",
params.file_path.display()
)),
let bytes = std::fs::read(&path).map_err(|e| match e.kind() {
std::io::ErrorKind::NotFound => {
ToolError::ExecutionFailed(format!("record not found: {}", path.display()))
}
_ => ToolError::ExecutionFailed(format!("read failed at {}: {e}", path.display())),
})?;
let text = String::from_utf8_lossy(&bytes).into_owned();
@@ -84,13 +68,13 @@ impl Tool for ReadTool {
offset + 1,
offset + rendered.line_count,
rendered.total_lines,
params.file_path.display()
path.display()
)
} else {
format!(
"Read {} line(s) from {}",
rendered.line_count,
params.file_path.display()
path.display()
)
};
@@ -157,14 +141,14 @@ mod tests {
}
#[tokio::test]
async fn read_returns_numbered_lines() {
async fn read_decision_by_slug() {
let (dir, layout) = setup();
let path = dir.path().join("memory/decisions/foo.md");
std::fs::create_dir_all(path.parent().unwrap()).unwrap();
std::fs::write(&path, "alpha\nbeta\n").unwrap();
let (_meta, tool) = read_tool(layout)();
let inp = serde_json::json!({ "file_path": path.to_str().unwrap() });
let inp = serde_json::json!({ "kind": "decision", "slug": "foo" });
let out = tool.execute(&inp.to_string()).await.unwrap();
let body = out.content.unwrap();
assert!(body.contains(" 1\talpha"));
@@ -172,24 +156,64 @@ mod tests {
}
#[tokio::test]
async fn rejects_outside_memory_tree() {
async fn read_summary_without_slug() {
let (dir, layout) = setup();
let other = dir.path().join("src/main.rs");
std::fs::create_dir_all(other.parent().unwrap()).unwrap();
std::fs::write(&other, "fn main() {}").unwrap();
let path = dir.path().join("memory/summary.md");
std::fs::create_dir_all(path.parent().unwrap()).unwrap();
std::fs::write(&path, "summary body\n").unwrap();
let (_, tool) = read_tool(layout)();
let inp = serde_json::json!({ "file_path": other.to_str().unwrap() });
let inp = serde_json::json!({ "kind": "summary" });
let out = tool.execute(&inp.to_string()).await.unwrap();
assert!(out.content.unwrap().contains("summary body"));
}
#[tokio::test]
async fn summary_with_slug_rejected() {
let (_dir, layout) = setup();
let (_, tool) = read_tool(layout)();
let inp = serde_json::json!({ "kind": "summary", "slug": "x" });
let err = tool.execute(&inp.to_string()).await.unwrap_err();
assert!(matches!(err, ToolError::InvalidArgument(_)));
}
#[tokio::test]
async fn rejects_relative_path() {
async fn decision_without_slug_rejected() {
let (_dir, layout) = setup();
let (_, tool) = read_tool(layout)();
let inp = serde_json::json!({ "file_path": "memory/summary.md" });
let inp = serde_json::json!({ "kind": "decision" });
let err = tool.execute(&inp.to_string()).await.unwrap_err();
assert!(matches!(err, ToolError::InvalidArgument(_)));
}
#[tokio::test]
async fn invalid_slug_rejected() {
let (_dir, layout) = setup();
let (_, tool) = read_tool(layout)();
let inp = serde_json::json!({ "kind": "decision", "slug": "Bad-Slug" });
let err = tool.execute(&inp.to_string()).await.unwrap_err();
assert!(matches!(err, ToolError::InvalidArgument(_)));
}
#[tokio::test]
async fn knowledge_path_resolution() {
let (dir, layout) = setup();
let path = dir.path().join("knowledge/policy.md");
std::fs::create_dir_all(path.parent().unwrap()).unwrap();
std::fs::write(&path, "k\n").unwrap();
let (_, tool) = read_tool(layout)();
let inp = serde_json::json!({ "kind": "knowledge", "slug": "policy" });
let out = tool.execute(&inp.to_string()).await.unwrap();
assert!(out.content.unwrap().contains("k"));
}
#[tokio::test]
async fn missing_file_returns_execution_failed() {
let (_dir, layout) = setup();
let (_, tool) = read_tool(layout)();
let inp = serde_json::json!({ "kind": "decision", "slug": "missing" });
let err = tool.execute(&inp.to_string()).await.unwrap_err();
assert!(matches!(err, ToolError::ExecutionFailed(_)));
}
}
crates/memory/src/tool/search.rs
+560
View File
@@ -0,0 +1,560 @@
//! `MemorySearch` / `KnowledgeSearch` tools.
//!
//! Both perform a case-insensitive substring scan over markdown record
//! files, returning a list of `{slug, kind, ..., excerpt}` entries.
//! Excerpts are `excerpt_lines` lines before and after the matched
//! line (so 2N+1 lines per excerpt when not clipped).
//!
//! - `MemorySearch` walks `memory/summary.md`, `memory/decisions/`,
//! `memory/requests/`. `memory/workflow/` and `memory/_staging/`
//! are excluded by construction.
//! - `KnowledgeSearch` walks `knowledge/*.md` and supports a `kind`
//! filter against the Knowledge frontmatter's `kind` field.
//!
//! No derived index — the file tree is the source of truth and is
//! re-scanned per call. grep 出現順: within a file by line order,
//! across files by sorted filename.
use std::path::{Path, PathBuf};
use std::sync::Arc;
use async_trait::async_trait;
use llm_worker::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use serde::{Deserialize, Serialize};
use crate::schema::{KnowledgeFrontmatter, split_frontmatter};
use crate::workspace::WorkspaceLayout;
const DEFAULT_HIT_LIMIT: usize = 20;
const DEFAULT_EXCERPT_LINES: usize = 3;
const MEMORY_SEARCH_DESCRIPTION: &str = "Search memory records (summary / decisions / \
requests) for a substring. Returns up to `hit_limit` matches as `{slug, kind, excerpt}` \
entries with line context. Use the returned `slug` + `kind` with MemoryRead to fetch \
the full record. Workflow and staging directories are not searched.";
const KNOWLEDGE_SEARCH_DESCRIPTION: &str = "Search knowledge records for a substring. \
Optional `kind` filters by the Knowledge frontmatter's `kind` field. Returns up to \
`hit_limit` matches as `{slug, kind, description, model_invokation, excerpt}` entries \
with line context. Use the returned `slug` with MemoryRead (kind=knowledge) for the \
full record.";
/// Tunables passed in from the manifest.
#[derive(Debug, Clone, Copy)]
pub struct SearchConfig {
pub hit_limit: usize,
/// Lines of context before and after each matched line.
pub excerpt_lines: usize,
}
impl Default for SearchConfig {
fn default() -> Self {
Self {
hit_limit: DEFAULT_HIT_LIMIT,
excerpt_lines: DEFAULT_EXCERPT_LINES,
}
}
}
#[derive(Debug, Deserialize, schemars::JsonSchema)]
struct MemorySearchParams {
/// Substring to search for. Case-insensitive.
query: String,
}
#[derive(Debug, Deserialize, schemars::JsonSchema)]
struct KnowledgeSearchParams {
/// Substring to search for. Case-insensitive.
query: String,
/// Optional filter on the Knowledge frontmatter's `kind` field.
#[serde(default)]
kind: Option<String>,
}
#[derive(Debug, Serialize)]
struct MemoryHit {
slug: String,
kind: &'static str,
excerpt: String,
}
#[derive(Debug, Serialize)]
struct KnowledgeHit {
slug: String,
kind: Option<String>,
description: Option<String>,
model_invokation: Option<bool>,
excerpt: String,
}
struct MemorySearchTool {
layout: WorkspaceLayout,
config: SearchConfig,
}
struct KnowledgeSearchTool {
layout: WorkspaceLayout,
config: SearchConfig,
}
#[async_trait]
impl Tool for MemorySearchTool {
async fn execute(&self, input_json: &str) -> Result<ToolOutput, ToolError> {
let params: MemorySearchParams = serde_json::from_str(input_json).map_err(|e| {
ToolError::InvalidArgument(format!("invalid MemorySearch input: {e}"))
})?;
let needle = validate_query(&params.query)?;
let mut hits: Vec<MemoryHit> = Vec::new();
let limit = self.config.hit_limit;
let ctx = self.config.excerpt_lines;
// summary
let summary_path = self.layout.summary_path();
if summary_path.is_file() {
scan_file(&summary_path, &needle, ctx, limit - hits.len(), |excerpt| {
hits.push(MemoryHit {
slug: "summary".to_string(),
kind: "summary",
excerpt,
});
});
}
// decisions
if hits.len() < limit {
for (path, slug) in list_md_files(&self.layout.decisions_dir()) {
if hits.len() >= limit {
break;
}
scan_file(&path, &needle, ctx, limit - hits.len(), |excerpt| {
hits.push(MemoryHit {
slug: slug.clone(),
kind: "decision",
excerpt,
});
});
}
}
// requests
if hits.len() < limit {
for (path, slug) in list_md_files(&self.layout.requests_dir()) {
if hits.len() >= limit {
break;
}
scan_file(&path, &needle, ctx, limit - hits.len(), |excerpt| {
hits.push(MemoryHit {
slug: slug.clone(),
kind: "request",
excerpt,
});
});
}
}
let body = serde_json::to_string_pretty(&hits)
.map_err(|e| ToolError::ExecutionFailed(format!("serialize hits: {e}")))?;
Ok(ToolOutput {
summary: format!("{} hit(s) for {:?}", hits.len(), params.query),
content: Some(body),
})
}
}
#[async_trait]
impl Tool for KnowledgeSearchTool {
async fn execute(&self, input_json: &str) -> Result<ToolOutput, ToolError> {
let params: KnowledgeSearchParams = serde_json::from_str(input_json).map_err(|e| {
ToolError::InvalidArgument(format!("invalid KnowledgeSearch input: {e}"))
})?;
let needle = validate_query(&params.query)?;
let kind_filter = params.kind.as_deref();
let mut hits: Vec<KnowledgeHit> = Vec::new();
let limit = self.config.hit_limit;
let ctx = self.config.excerpt_lines;
for (path, slug) in list_md_files(&self.layout.knowledge_dir()) {
if hits.len() >= limit {
break;
}
// Try to parse frontmatter for description/model_invokation/kind.
let raw = match std::fs::read_to_string(&path) {
Ok(s) => s,
Err(_) => continue,
};
let fm = parse_knowledge_frontmatter(&raw);
// kind filter applies to the frontmatter's kind field.
if let Some(filter) = kind_filter {
let matches = fm
.as_ref()
.map(|f| f.kind.as_str() == filter)
.unwrap_or(false);
if !matches {
continue;
}
}
let kind = fm.as_ref().map(|f| f.kind.clone());
let description = fm.as_ref().map(|f| f.description.clone());
let model_invokation = fm.as_ref().map(|f| f.model_invokation);
scan_text(&raw, &needle, ctx, limit - hits.len(), |excerpt| {
hits.push(KnowledgeHit {
slug: slug.clone(),
kind: kind.clone(),
description: description.clone(),
model_invokation,
excerpt,
});
});
}
let body = serde_json::to_string_pretty(&hits)
.map_err(|e| ToolError::ExecutionFailed(format!("serialize hits: {e}")))?;
Ok(ToolOutput {
summary: format!("{} hit(s) for {:?}", hits.len(), params.query),
content: Some(body),
})
}
}
fn validate_query(query: &str) -> Result<String, ToolError> {
if query.trim().is_empty() {
return Err(ToolError::InvalidArgument(
"query must not be empty".into(),
));
}
Ok(query.to_lowercase())
}
/// Sorted list of `(path, slug)` for `*.md` files directly under `dir`.
/// Returns empty if the directory doesn't exist.
fn list_md_files(dir: &Path) -> Vec<(PathBuf, String)> {
let mut out: Vec<(PathBuf, String)> = Vec::new();
let entries = match std::fs::read_dir(dir) {
Ok(it) => it,
Err(_) => return out,
};
for entry in entries.flatten() {
let path = entry.path();
if !path.is_file() {
continue;
}
let name = match path.file_name().and_then(|n| n.to_str()) {
Some(n) => n,
None => continue,
};
let slug = match name.strip_suffix(".md") {
Some(s) => s.to_string(),
None => continue,
};
out.push((path, slug));
}
out.sort_by(|a, b| a.1.cmp(&b.1));
out
}
fn scan_file(
path: &Path,
needle_lower: &str,
ctx: usize,
remaining: usize,
mut on_match: impl FnMut(String),
) {
if remaining == 0 {
return;
}
let text = match std::fs::read_to_string(path) {
Ok(t) => t,
Err(_) => return,
};
scan_text(&text, needle_lower, ctx, remaining, |e| on_match(e));
}
fn scan_text(
text: &str,
needle_lower: &str,
ctx: usize,
remaining: usize,
mut on_match: impl FnMut(String),
) {
if remaining == 0 {
return;
}
let lines: Vec<&str> = text.lines().collect();
let mut produced = 0;
for (i, line) in lines.iter().enumerate() {
if produced >= remaining {
break;
}
if line.to_lowercase().contains(needle_lower) {
let start = i.saturating_sub(ctx);
let end = i.saturating_add(ctx + 1).min(lines.len());
let excerpt = lines[start..end].join("\n");
on_match(excerpt);
produced += 1;
}
}
}
/// Best-effort frontmatter parse. Returns `None` if missing/malformed
/// — search still finds matches in the body even when the header is
/// broken.
fn parse_knowledge_frontmatter(raw: &str) -> Option<KnowledgeFrontmatter> {
let (yaml, _body) = split_frontmatter(raw).ok()?;
serde_yaml::from_str::<KnowledgeFrontmatter>(yaml).ok()
}
pub fn memory_search_tool(layout: WorkspaceLayout, config: SearchConfig) -> ToolDefinition {
Arc::new(move || {
let schema = schemars::schema_for!(MemorySearchParams);
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
let meta = ToolMeta::new("MemorySearch")
.description(MEMORY_SEARCH_DESCRIPTION)
.input_schema(schema_value);
let tool: Arc<dyn Tool> = Arc::new(MemorySearchTool {
layout: layout.clone(),
config,
});
(meta, tool)
})
}
pub fn knowledge_search_tool(layout: WorkspaceLayout, config: SearchConfig) -> ToolDefinition {
Arc::new(move || {
let schema = schemars::schema_for!(KnowledgeSearchParams);
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
let meta = ToolMeta::new("KnowledgeSearch")
.description(KNOWLEDGE_SEARCH_DESCRIPTION)
.input_schema(schema_value);
let tool: Arc<dyn Tool> = Arc::new(KnowledgeSearchTool {
layout: layout.clone(),
config,
});
(meta, tool)
})
}
#[cfg(test)]
mod tests {
use super::*;
use chrono::Utc;
use tempfile::TempDir;
fn now() -> String {
Utc::now().to_rfc3339()
}
fn setup() -> (TempDir, WorkspaceLayout) {
let dir = TempDir::new().unwrap();
let layout = WorkspaceLayout::new(dir.path().to_path_buf());
std::fs::create_dir_all(dir.path().join("memory/decisions")).unwrap();
std::fs::create_dir_all(dir.path().join("memory/requests")).unwrap();
std::fs::create_dir_all(dir.path().join("memory/workflow")).unwrap();
std::fs::create_dir_all(dir.path().join("memory/_staging")).unwrap();
std::fs::create_dir_all(dir.path().join("knowledge")).unwrap();
(dir, layout)
}
fn write_decision(dir: &Path, slug: &str, body: &str) {
let path = dir.join("memory/decisions").join(format!("{slug}.md"));
let content = format!(
"---\ncreated_at: {n}\nupdated_at: {n}\nsources: []\nstatus: open\n---\n{body}",
n = now()
);
std::fs::write(path, content).unwrap();
}
fn write_knowledge(dir: &Path, slug: &str, kind: &str, description: &str, body: &str) {
let path = dir.join("knowledge").join(format!("{slug}.md"));
let content = format!(
"---\ncreated_at: {n}\nupdated_at: {n}\nkind: {kind}\ndescription: \"{description}\"\nmodel_invokation: false\nuser_invocable: true\nlast_sources: []\n---\n{body}",
n = now()
);
std::fs::write(path, content).unwrap();
}
fn parse_hits<T: for<'de> serde::Deserialize<'de>>(out: &ToolOutput) -> Vec<T> {
serde_json::from_str(out.content.as_ref().unwrap()).unwrap()
}
#[derive(Deserialize)]
struct OwnedMemoryHit {
slug: String,
kind: String,
excerpt: String,
}
#[derive(Deserialize)]
struct OwnedKnowledgeHit {
slug: String,
kind: Option<String>,
description: Option<String>,
model_invokation: Option<bool>,
excerpt: String,
}
#[tokio::test]
async fn memory_search_finds_decision_body() {
let (dir, layout) = setup();
write_decision(dir.path(), "alpha", "we chose Ollama because it works\n");
write_decision(dir.path(), "beta", "no match here\n");
let (_, tool) = memory_search_tool(layout, SearchConfig::default())();
let inp = serde_json::json!({ "query": "ollama" });
let out = tool.execute(&inp.to_string()).await.unwrap();
let hits: Vec<OwnedMemoryHit> = parse_hits(&out);
assert_eq!(hits.len(), 1);
assert_eq!(hits[0].slug, "alpha");
assert_eq!(hits[0].kind, "decision");
assert!(hits[0].excerpt.to_lowercase().contains("ollama"));
}
#[tokio::test]
async fn memory_search_finds_summary() {
let (dir, layout) = setup();
let summary_path = dir.path().join("memory/summary.md");
std::fs::write(
&summary_path,
format!("---\nupdated_at: {n}\n---\nthe needle is here\n", n = now()),
)
.unwrap();
let (_, tool) = memory_search_tool(layout, SearchConfig::default())();
let inp = serde_json::json!({ "query": "needle" });
let out = tool.execute(&inp.to_string()).await.unwrap();
let hits: Vec<OwnedMemoryHit> = parse_hits(&out);
assert_eq!(hits.len(), 1);
assert_eq!(hits[0].slug, "summary");
assert_eq!(hits[0].kind, "summary");
}
#[tokio::test]
async fn memory_search_excludes_workflow_and_staging() {
let (dir, layout) = setup();
// Workflow and staging files contain the needle but must be ignored.
let wf = dir.path().join("memory/workflow/wf.md");
std::fs::write(&wf, "needle in workflow\n").unwrap();
let stg = dir.path().join("memory/_staging/abc.json");
std::fs::write(&stg, "needle in staging\n").unwrap();
let (_, tool) = memory_search_tool(layout, SearchConfig::default())();
let inp = serde_json::json!({ "query": "needle" });
let out = tool.execute(&inp.to_string()).await.unwrap();
let hits: Vec<OwnedMemoryHit> = parse_hits(&out);
assert!(hits.is_empty(), "got hits: {:?}", out.content);
}
#[tokio::test]
async fn memory_search_respects_hit_limit() {
let (dir, layout) = setup();
for i in 0..10 {
write_decision(dir.path(), &format!("rec-{i}"), "needle line\n");
}
let cfg = SearchConfig {
hit_limit: 3,
excerpt_lines: 1,
};
let (_, tool) = memory_search_tool(layout, cfg)();
let inp = serde_json::json!({ "query": "needle" });
let out = tool.execute(&inp.to_string()).await.unwrap();
let hits: Vec<OwnedMemoryHit> = parse_hits(&out);
assert_eq!(hits.len(), 3);
}
#[tokio::test]
async fn memory_search_excerpt_includes_context_lines() {
let (dir, layout) = setup();
write_decision(
dir.path(),
"ctx",
"line a\nline b\nNEEDLE here\nline d\nline e\n",
);
let cfg = SearchConfig {
hit_limit: 5,
excerpt_lines: 1,
};
let (_, tool) = memory_search_tool(layout, cfg)();
let inp = serde_json::json!({ "query": "needle" });
let out = tool.execute(&inp.to_string()).await.unwrap();
let hits: Vec<OwnedMemoryHit> = parse_hits(&out);
assert_eq!(hits.len(), 1);
let e = &hits[0].excerpt;
assert!(e.contains("line b"));
assert!(e.contains("NEEDLE here"));
assert!(e.contains("line d"));
assert!(!e.contains("line a"));
assert!(!e.contains("line e"));
}
#[tokio::test]
async fn memory_search_empty_query_rejected() {
let (_dir, layout) = setup();
let (_, tool) = memory_search_tool(layout, SearchConfig::default())();
let inp = serde_json::json!({ "query": " " });
let err = tool.execute(&inp.to_string()).await.unwrap_err();
assert!(matches!(err, ToolError::InvalidArgument(_)));
}
#[tokio::test]
async fn knowledge_search_returns_frontmatter_fields() {
let (dir, layout) = setup();
write_knowledge(
dir.path(),
"policy",
"policy",
"the policy doc",
"Ollama first\n",
);
let (_, tool) = knowledge_search_tool(layout, SearchConfig::default())();
let inp = serde_json::json!({ "query": "ollama" });
let out = tool.execute(&inp.to_string()).await.unwrap();
let hits: Vec<OwnedKnowledgeHit> = parse_hits(&out);
assert_eq!(hits.len(), 1);
assert_eq!(hits[0].slug, "policy");
assert_eq!(hits[0].kind.as_deref(), Some("policy"));
assert_eq!(hits[0].description.as_deref(), Some("the policy doc"));
assert_eq!(hits[0].model_invokation, Some(false));
assert!(hits[0].excerpt.to_lowercase().contains("ollama"));
}
#[tokio::test]
async fn knowledge_search_kind_filter() {
let (dir, layout) = setup();
write_knowledge(dir.path(), "p1", "policy", "d1", "needle\n");
write_knowledge(dir.path(), "h1", "howto", "d2", "needle\n");
let (_, tool) = knowledge_search_tool(layout, SearchConfig::default())();
let inp = serde_json::json!({ "query": "needle", "kind": "howto" });
let out = tool.execute(&inp.to_string()).await.unwrap();
let hits: Vec<OwnedKnowledgeHit> = parse_hits(&out);
assert_eq!(hits.len(), 1);
assert_eq!(hits[0].slug, "h1");
}
#[tokio::test]
async fn knowledge_search_searches_frontmatter_too() {
// Spec completion criteria: "frontmatter 含む全文から excerpt 付きでヒットが返る"
let (dir, layout) = setup();
write_knowledge(dir.path(), "p", "policy", "mentions xyzzy here", "body\n");
let (_, tool) = knowledge_search_tool(layout, SearchConfig::default())();
let inp = serde_json::json!({ "query": "xyzzy" });
let out = tool.execute(&inp.to_string()).await.unwrap();
let hits: Vec<OwnedKnowledgeHit> = parse_hits(&out);
assert_eq!(hits.len(), 1);
assert_eq!(hits[0].slug, "p");
}
#[tokio::test]
async fn knowledge_search_no_matches_returns_empty() {
let (dir, layout) = setup();
write_knowledge(dir.path(), "p", "policy", "d", "no match\n");
let (_, tool) = knowledge_search_tool(layout, SearchConfig::default())();
let inp = serde_json::json!({ "query": "absent" });
let out = tool.execute(&inp.to_string()).await.unwrap();
let hits: Vec<OwnedKnowledgeHit> = parse_hits(&out);
assert!(hits.is_empty());
}
}
crates/memory/src/tool/write.rs
+46 -53
View File
@@ -1,12 +1,11 @@
//! `MemoryWrite` tool.
//!
//! Creates or overwrites a memory or knowledge record with full content.
//! Creates or overwrites a memory or knowledge record by `(kind, slug)`.
//! Pre-write Linter validates frontmatter, slug uniqueness (Create only),
//! reference integrity, size limits, and the workflow-write ban. On any
//! Linter error the tool returns `ToolError::InvalidArgument` with all
//! violations aggregated and the file is **not** written.
use std::path::PathBuf;
use std::sync::Arc;
use async_trait::async_trait;
@@ -14,22 +13,27 @@ use llm_worker::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use serde::Deserialize;
use crate::linter::{LintReport, Linter, WriteMode};
use crate::tool::MemoryToolKind;
use crate::workspace::WorkspaceLayout;
const DESCRIPTION: &str = "Create or overwrite a memory or knowledge record file. \
Path must be absolute and lie inside the workspace's `memory/` or `knowledge/` \
tree. Frontmatter is validated before the file is written; on validation \
failure no write occurs and every violation is returned in the error message.";
const DESCRIPTION: &str = "Create or overwrite a memory or knowledge record by \
`kind` + `slug`. `kind`: summary | decision | request | knowledge. For `summary` \
omit `slug`. Frontmatter is validated before write; on validation failure no \
write occurs and every violation is returned in the error message.";
#[derive(Debug, Deserialize, schemars::JsonSchema)]
struct WriteParams {
/// Absolute path under the workspace's `memory/` or `knowledge/` tree.
file_path: PathBuf,
/// Record kind: `summary` | `decision` | `request` | `knowledge`.
kind: MemoryToolKind,
/// Slug. Required for everything except `summary`; forbidden for `summary`.
#[serde(default)]
slug: Option<String>,
/// Full file contents (frontmatter + body).
content: String,
}
struct WriteTool {
layout: WorkspaceLayout,
linter: Linter,
}
@@ -40,26 +44,21 @@ impl Tool for WriteTool {
ToolError::InvalidArgument(format!("invalid MemoryWrite input: {e}"))
})?;
if !params.file_path.is_absolute() {
return Err(ToolError::InvalidArgument(format!(
"file_path must be absolute: {}",
params.file_path.display()
)));
}
let path = params.kind.resolve_path(&self.layout, params.slug.as_deref())?;
let already_exists = params.file_path.exists();
let already_exists = path.exists();
let mode = if already_exists {
WriteMode::Update
} else {
WriteMode::Create
};
let report = self.linter.lint(&params.file_path, &params.content, mode);
let report = self.linter.lint(&path, &params.content, mode);
if report.has_errors() {
return Err(ToolError::InvalidArgument(format_report(&report)));
}
if let Some(parent) = params.file_path.parent() {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent).map_err(|e| {
ToolError::ExecutionFailed(format!(
"failed to create directory {}: {e}",
@@ -67,17 +66,14 @@ impl Tool for WriteTool {
))
})?;
}
std::fs::write(&params.file_path, params.content.as_bytes()).map_err(|e| {
ToolError::ExecutionFailed(format!(
"failed to write {}: {e}",
params.file_path.display()
))
std::fs::write(&path, params.content.as_bytes()).map_err(|e| {
ToolError::ExecutionFailed(format!("failed to write {}: {e}", path.display()))
})?;
let summary = format!(
"{} {}{}",
if already_exists { "Overwrote" } else { "Created" },
params.file_path.display(),
path.display(),
warning_tail(&report),
);
Ok(ToolOutput {
@@ -122,6 +118,7 @@ pub fn write_tool(layout: WorkspaceLayout) -> ToolDefinition {
.description(DESCRIPTION)
.input_schema(schema_value);
let tool: Arc<dyn Tool> = Arc::new(WriteTool {
layout: layout.clone(),
linter: Linter::new(layout.clone()),
});
(meta, tool)
@@ -154,7 +151,7 @@ mod tests {
assert_eq!(meta.name, "MemoryWrite");
let inp = serde_json::json!({
"file_path": path.to_str().unwrap(),
"kind": "summary",
"content": content,
});
let out = tool.execute(&inp.to_string()).await.unwrap();
@@ -162,30 +159,10 @@ mod tests {
assert!(path.exists());
}
#[tokio::test]
async fn write_rejects_workflow() {
let (dir, layout) = setup();
let path = dir.path().join("memory/workflow/wf.md");
let content = format!(
"---\nupdated_at: {n}\ndescription: x\nauto_invoke: false\nuser_invocable: true\n---\n",
n = now()
);
let (_, tool) = write_tool(layout)();
let inp = serde_json::json!({
"file_path": path.to_str().unwrap(),
"content": content,
});
let err = tool.execute(&inp.to_string()).await.unwrap_err();
let msg = format!("{err}");
assert!(msg.contains("workflow"), "unexpected error: {msg}");
assert!(!path.exists(), "workflow file must not be written");
}
#[tokio::test]
async fn write_aggregates_multiple_errors() {
let (dir, layout) = setup();
let path = dir.path().join("memory/decisions/foo.md");
// Missing required `status` field AND body too long.
let (_dir, layout) = setup();
// Missing required `status` field for decisions.
let huge = "x".repeat(8001);
let content = format!(
"---\ncreated_at: {n}\nupdated_at: {n}\nsources: []\n---\n{huge}",
@@ -193,7 +170,8 @@ mod tests {
);
let (_, tool) = write_tool(layout)();
let inp = serde_json::json!({
"file_path": path.to_str().unwrap(),
"kind": "decision",
"slug": "foo",
"content": content,
});
let err = tool.execute(&inp.to_string()).await.unwrap_err();
@@ -202,7 +180,7 @@ mod tests {
}
#[tokio::test]
async fn write_blocks_create_when_existing() {
async fn write_update_existing() {
let (dir, layout) = setup();
let path = dir.path().join("memory/decisions/foo.md");
std::fs::create_dir_all(path.parent().unwrap()).unwrap();
@@ -212,10 +190,10 @@ mod tests {
);
std::fs::write(&path, &initial).unwrap();
// Same content as a re-write should pass (Update mode).
let (_, tool) = write_tool(layout.clone())();
let inp = serde_json::json!({
"file_path": path.to_str().unwrap(),
"kind": "decision",
"slug": "foo",
"content": initial,
});
let out = tool.execute(&inp.to_string()).await.unwrap();
@@ -223,11 +201,11 @@ mod tests {
}
#[tokio::test]
async fn write_rejects_non_absolute() {
async fn write_decision_requires_slug() {
let (_dir, layout) = setup();
let (_, tool) = write_tool(layout)();
let inp = serde_json::json!({
"file_path": "memory/summary.md",
"kind": "decision",
"content": "ignored",
});
let err = tool.execute(&inp.to_string()).await.unwrap_err();
@@ -241,10 +219,25 @@ mod tests {
let bad = "no frontmatter at all";
let (_, tool) = write_tool(layout)();
let inp = serde_json::json!({
"file_path": path.to_str().unwrap(),
"kind": "decision",
"slug": "foo",
"content": bad,
});
assert!(tool.execute(&inp.to_string()).await.is_err());
assert!(!path.exists());
}
#[tokio::test]
async fn workflow_kind_not_acceptable() {
// The MemoryToolKind enum doesn't include Workflow, so deserialization fails.
let (_dir, layout) = setup();
let (_, tool) = write_tool(layout)();
let inp = serde_json::json!({
"kind": "workflow",
"slug": "wf",
"content": "---\n---\n",
});
let err = tool.execute(&inp.to_string()).await.unwrap_err();
assert!(matches!(err, ToolError::InvalidArgument(_)));
}
}