feat(primitives): 10 Rust crates extracted from LBM (Genesis-scrubbed)

- kei-router — keyword-dispatch meta-tool (CfC ML fallback removed)
- kei-sage — Obsidian-style knowledge graph, FTS5 + BFS + PageRank
- kei-task — task DAG with deps, milestones, dependency-chain queries
- kei-chat-store — Claude conversation session persistence + FTS search
- kei-crossdomain — typed-edge store + BFS cross-domain glue
- kei-search-core — 3-wave deep research with microcent budget cap
- kei-content-store — asset + prompt + campaign registry
- kei-social-store — people + interactions CRM (lite)
- kei-curator — edge-decay graph hygiene utility
- kei-auth — multi-tenant session tokens (replaces single-bearer)

Genesis-scan pre-import pass: skipped pkg/mxl1/*, pkg/inference/*, pkg/trainer/*,
pkg/nc01/*, internal/ml/* (all Genesis/CfC adjacent, sensitive IP).
Security: skipped tools_threat/radio/protocol/med/mlreg (offensive/banned).
Domain verticals skipped: hr/legal/infra/ops/api/osint/edu/geo/hw/finance.

New 'mcp' profile in MANIFEST.toml bundles all 10 for MCP server deployment.

Workspace now 24 crates, cargo check --workspace clean, 94 workspace tests pass.
This commit is contained in:
Parfii-bot 2026-04-22 12:48:56 +08:00
parent cab78d68f7
commit adc007b7b0
98 changed files with 5089 additions and 1 deletions

View file

@ -21,7 +21,8 @@ core = ["tomd", "genesis-scan"]
frontend = ["mock-render", "visual-diff", "tokens-sync", "design-scrape", "live-preview", "figma-tokens", "frontend-inspect", "screenshot-decode"]
ops = ["kei-ledger", "ssh-check", "firewall-diff", "provision-hetzner", "provision-vultr", "harden-base", "metrics-scrape", "log-ship"]
dev = ["kei-migrate", "kei-changelog", "kei-ci-lint", "kei-docs-scaffold", "kei-memory", "kei-conflict-scan", "kei-refactor-engine", "kei-graph-check", "kei-store"]
full = ["tomd", "genesis-scan", "kei-ledger", "kei-migrate", "kei-changelog", "ssh-check", "firewall-diff", "mock-render", "visual-diff", "tokens-sync", "design-scrape", "live-preview", "figma-tokens", "frontend-inspect", "screenshot-decode", "provision-hetzner", "provision-vultr", "harden-base", "metrics-scrape", "log-ship", "kei-ci-lint", "kei-docs-scaffold", "kei-memory", "kei-conflict-scan", "kei-refactor-engine", "kei-graph-check", "kei-store"]
mcp = ["kei-router", "kei-sage", "kei-task", "kei-chat-store", "kei-crossdomain", "kei-search-core", "kei-content-store", "kei-social-store", "kei-curator", "kei-auth"]
full = ["tomd", "genesis-scan", "kei-ledger", "kei-migrate", "kei-changelog", "ssh-check", "firewall-diff", "mock-render", "visual-diff", "tokens-sync", "design-scrape", "live-preview", "figma-tokens", "frontend-inspect", "screenshot-decode", "provision-hetzner", "provision-vultr", "harden-base", "metrics-scrape", "log-ship", "kei-ci-lint", "kei-docs-scaffold", "kei-memory", "kei-conflict-scan", "kei-refactor-engine", "kei-graph-check", "kei-store", "kei-router", "kei-sage", "kei-task", "kei-chat-store", "kei-crossdomain", "kei-search-core", "kei-content-store", "kei-social-store", "kei-curator", "kei-auth"]
# --- shell primitives (13) -------------------------------------------------
@ -188,3 +189,65 @@ kind = "rust"
crate = "kei-store"
deps = ["git2 (vendored libgit2)"]
desc = "Memory-repo backend abstraction — GitHub / Forgejo / Gitea / Filesystem / S3 (S3 = MVP stub)"
# --- v0.14 LBM port (10) ---------------------------------------------------
[primitive.kei-router]
kind = "rust"
crate = "kei-router"
deps = ["regex"]
desc = "Natural-language query → tool-call router (LBM pkg/keirouter port, no ML)"
[primitive.kei-sage]
kind = "rust"
crate = "kei-sage"
deps = ["rusqlite bundled (FTS5 enabled)"]
desc = "Obsidian-style knowledge graph with FTS5, BFS, PageRank (LBM internal/sage port)"
[primitive.kei-task]
kind = "rust"
crate = "kei-task"
deps = ["rusqlite bundled (FTS5 enabled)"]
desc = "Task DAG + deps + milestones (LBM internal/task port)"
[primitive.kei-chat-store]
kind = "rust"
crate = "kei-chat-store"
deps = ["rusqlite bundled (FTS5 enabled)"]
desc = "Session persistence for Claude chats (LBM internal/chat port)"
[primitive.kei-crossdomain]
kind = "rust"
crate = "kei-crossdomain"
deps = ["rusqlite bundled"]
desc = "Cross-domain typed-edge store + BFS + auto-link (LBM internal/crossdomain port)"
[primitive.kei-search-core]
kind = "rust"
crate = "kei-search-core"
deps = ["rusqlite bundled"]
desc = "3-wave research engine with budget cap; fetch interface frozen (LBM internal/search port)"
[primitive.kei-content-store]
kind = "rust"
crate = "kei-content-store"
deps = ["rusqlite bundled", "sha2"]
desc = "Asset + prompt + campaign registry (LBM internal/content port)"
[primitive.kei-social-store]
kind = "rust"
crate = "kei-social-store"
deps = ["rusqlite bundled (FTS5 enabled)"]
desc = "People + interaction CRM lite (LBM internal/social port)"
[primitive.kei-curator]
kind = "rust"
crate = "kei-curator"
deps = ["rusqlite bundled"]
desc = "Edge decay + orphan prune for cross-domain graphs (LBM internal/curator port)"
[primitive.kei-auth]
kind = "rust"
crate = "kei-auth"
deps = ["rusqlite bundled", "hmac", "sha2"]
desc = "Multi-tenant session tokens with scopes + HMAC-signed expiry (rewrite, not port)"

View file

@ -909,6 +909,23 @@ dependencies = [
"wasm-bindgen",
]
[[package]]
name = "kei-auth"
version = "0.1.0"
dependencies = [
"anyhow",
"base64 0.22.1",
"chrono",
"clap",
"hmac",
"rand",
"rusqlite",
"serde",
"serde_json",
"sha2",
"tempfile",
]
[[package]]
name = "kei-changelog"
version = "0.1.0"
@ -920,6 +937,20 @@ dependencies = [
"regex",
]
[[package]]
name = "kei-chat-store"
version = "0.1.0"
dependencies = [
"anyhow",
"chrono",
"clap",
"rusqlite",
"serde",
"serde_json",
"tempfile",
"uuid",
]
[[package]]
name = "kei-conflict-scan"
version = "0.1.0"
@ -933,6 +964,46 @@ dependencies = [
"walkdir",
]
[[package]]
name = "kei-content-store"
version = "0.1.0"
dependencies = [
"anyhow",
"chrono",
"clap",
"rusqlite",
"serde",
"serde_json",
"sha2",
"tempfile",
]
[[package]]
name = "kei-crossdomain"
version = "0.1.0"
dependencies = [
"anyhow",
"chrono",
"clap",
"rusqlite",
"serde",
"serde_json",
"tempfile",
]
[[package]]
name = "kei-curator"
version = "0.1.0"
dependencies = [
"anyhow",
"chrono",
"clap",
"rusqlite",
"serde",
"serde_json",
"tempfile",
]
[[package]]
name = "kei-graph-check"
version = "0.1.0"
@ -996,6 +1067,56 @@ dependencies = [
"tempfile",
]
[[package]]
name = "kei-router"
version = "0.1.0"
dependencies = [
"anyhow",
"clap",
"regex",
"serde",
"serde_json",
]
[[package]]
name = "kei-sage"
version = "0.1.0"
dependencies = [
"anyhow",
"chrono",
"clap",
"rusqlite",
"serde",
"serde_json",
"tempfile",
]
[[package]]
name = "kei-search-core"
version = "0.1.0"
dependencies = [
"anyhow",
"chrono",
"clap",
"rusqlite",
"serde",
"serde_json",
"tempfile",
]
[[package]]
name = "kei-social-store"
version = "0.1.0"
dependencies = [
"anyhow",
"chrono",
"clap",
"rusqlite",
"serde",
"serde_json",
"tempfile",
]
[[package]]
name = "kei-store"
version = "0.1.0"
@ -1009,6 +1130,19 @@ dependencies = [
"toml",
]
[[package]]
name = "kei-task"
version = "0.1.0"
dependencies = [
"anyhow",
"chrono",
"clap",
"rusqlite",
"serde",
"serde_json",
"tempfile",
]
[[package]]
name = "lazy_static"
version = "1.5.0"
@ -2315,6 +2449,17 @@ version = "0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821"
[[package]]
name = "uuid"
version = "1.23.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ddd74a9687298c6858e9b88ec8935ec45d22e8fd5e6394fa1bd4e99a87789c76"
dependencies = [
"getrandom 0.4.2",
"js-sys",
"wasm-bindgen",
]
[[package]]
name = "vcpkg"
version = "0.2.15"

View file

@ -15,6 +15,17 @@ members = [
"kei-refactor-engine",
"kei-graph-check",
"kei-store",
# v0.14 LBM port — 10 new MCP-core primitives
"kei-router",
"kei-sage",
"kei-task",
"kei-chat-store",
"kei-crossdomain",
"kei-search-core",
"kei-content-store",
"kei-social-store",
"kei-curator",
"kei-auth",
]
[workspace.package]

View file

@ -0,0 +1,29 @@
[package]
name = "kei-auth"
version = "0.1.0"
edition = "2021"
rust-version = "1.75"
description = "Multi-tenant session tokens with scopes + HMAC-signed expiry (SQLite backend)."
[[bin]]
name = "kei-auth"
path = "src/main.rs"
[lib]
name = "kei_auth"
path = "src/lib.rs"
[dependencies]
rusqlite = { version = "0.31", features = ["bundled"] }
clap = { version = "4", features = ["derive"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
anyhow = "1"
chrono = { version = "0.4", default-features = false, features = ["clock"] }
hmac = "0.12"
sha2 = "0.10"
base64 = "0.22"
rand = "0.8"
[dev-dependencies]
tempfile = "3"

View file

@ -0,0 +1,25 @@
//! HMAC-SHA256 signer for token bodies.
use ::hmac::{Hmac, Mac};
use anyhow::{anyhow, Result};
use base64::Engine;
use sha2::Sha256;
type H = Hmac<Sha256>;
/// Sign `body` with `key`. Returns URL-safe base64 MAC.
pub fn sign(key: &[u8], body: &[u8]) -> String {
let mut mac = <H as Mac>::new_from_slice(key).expect("HMAC accepts any key size");
mac.update(body);
base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(mac.finalize().into_bytes())
}
/// Verify `body` against MAC. Returns Err if mismatch.
pub fn verify(key: &[u8], body: &[u8], mac_b64: &str) -> Result<()> {
let mut mac = <H as Mac>::new_from_slice(key).expect("HMAC accepts any key size");
mac.update(body);
let bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD
.decode(mac_b64)
.map_err(|e| anyhow!("bad b64 mac: {e}"))?;
mac.verify_slice(&bytes).map_err(|_| anyhow!("hmac mismatch"))
}

View file

@ -0,0 +1,15 @@
//! kei-auth — multi-tenant token auth. Replaces LBM's single LBM_MCP_TOKEN.
//!
//! Cubes:
//! - [`schema`] — SQLite tables for users + tokens
//! - [`hmac`] — HMAC-SHA256 signing helpers
//! - [`tokens`] — issue / verify / revoke / list
//! - [`scopes`] — read / write / admin enum + checks
pub mod hmac;
pub mod schema;
pub mod scopes;
pub mod tokens;
pub use scopes::Scope;
pub use tokens::{issue, revoke, verify, VerifyOutcome};

View file

@ -0,0 +1,70 @@
//! kei-auth CLI — issue/verify/revoke.
use clap::{Parser, Subcommand};
use kei_auth::schema::open;
use kei_auth::scopes::Scope;
use kei_auth::tokens::{issue, revoke, verify};
use std::path::PathBuf;
use std::process::ExitCode;
use std::str::FromStr;
#[derive(Parser)]
#[command(name = "kei-auth", version)]
struct Cli {
#[arg(long)] db: Option<PathBuf>,
/// HMAC signing key (env KEI_AUTH_KEY fallback).
#[arg(long)] key: Option<String>,
#[command(subcommand)] cmd: Cmd,
}
#[derive(Subcommand)]
enum Cmd {
Issue { #[arg(long)] user: String,
#[arg(long)] project: String,
#[arg(long, default_value = "read")] scope: String,
#[arg(long, default_value_t = 86400)] ttl: i64 },
Verify { token: String },
Revoke { token: String },
}
fn db_path(o: Option<PathBuf>) -> PathBuf {
if let Some(p) = o { return p; }
if let Ok(e) = std::env::var("KEI_AUTH_DB") { return PathBuf::from(e); }
let home = std::env::var("HOME").unwrap_or_else(|_| ".".into());
PathBuf::from(home).join(".claude/auth/auth.sqlite")
}
fn key(cli_key: Option<String>) -> anyhow::Result<Vec<u8>> {
if let Some(k) = cli_key { return Ok(k.into_bytes()); }
let k = std::env::var("KEI_AUTH_KEY")
.map_err(|_| anyhow::anyhow!("provide --key or set KEI_AUTH_KEY"))?;
Ok(k.into_bytes())
}
fn run() -> anyhow::Result<()> {
let cli = Cli::parse();
let conn = open(&db_path(cli.db))?;
let k = key(cli.key)?;
match cli.cmd {
Cmd::Issue { user, project, scope, ttl } => {
let sc = Scope::from_str(&scope).map_err(|e| anyhow::anyhow!(e))?;
println!("{}", issue(&conn, &user, &project, sc, ttl, &k)?);
}
Cmd::Verify { token } => {
let out = verify(&conn, &token, &k)?;
println!("user={} project={} scope={}", out.user_id, out.project, out.scope);
}
Cmd::Revoke { token } => {
let n = revoke(&conn, &token)?;
println!("revoked {} row(s)", n);
}
}
Ok(())
}
fn main() -> ExitCode {
match run() {
Ok(()) => ExitCode::SUCCESS,
Err(e) => { eprintln!("kei-auth: {e:#}"); ExitCode::from(1) }
}
}

View file

@ -0,0 +1,36 @@
use anyhow::{Context, Result};
use rusqlite::Connection;
use std::path::Path;
pub fn open(path: &Path) -> Result<Connection> {
if let Some(parent) = path.parent() {
let _ = std::fs::create_dir_all(parent);
}
let conn = Connection::open(path).context("open sqlite")?;
create_schema(&conn)?;
Ok(conn)
}
pub fn open_memory() -> Result<Connection> {
let conn = Connection::open_in_memory()?;
create_schema(&conn)?;
Ok(conn)
}
pub fn create_schema(conn: &Connection) -> Result<()> {
conn.execute_batch(r#"
CREATE TABLE IF NOT EXISTS auth_tokens (
id INTEGER PRIMARY KEY,
token_hash TEXT NOT NULL UNIQUE,
user_id TEXT NOT NULL,
project TEXT NOT NULL,
scope TEXT NOT NULL CHECK(scope IN ('read','write','admin')),
expires_at INTEGER NOT NULL,
created_at INTEGER NOT NULL,
revoked_at INTEGER DEFAULT 0
);
CREATE INDEX IF NOT EXISTS idx_tok_user ON auth_tokens(user_id);
CREATE INDEX IF NOT EXISTS idx_tok_project ON auth_tokens(project);
"#)?;
Ok(())
}

View file

@ -0,0 +1,45 @@
use serde::{Deserialize, Serialize};
use std::fmt;
use std::str::FromStr;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum Scope {
Read,
Write,
Admin,
}
impl Scope {
pub fn as_str(&self) -> &'static str {
match self { Scope::Read => "read", Scope::Write => "write", Scope::Admin => "admin" }
}
/// Admin ⊇ Write ⊇ Read.
pub fn allows(&self, required: Scope) -> bool {
use Scope::*;
match (self, required) {
(Admin, _) => true,
(Write, Read) | (Write, Write) => true,
(Read, Read) => true,
_ => false,
}
}
}
impl fmt::Display for Scope {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(self.as_str())
}
}
impl FromStr for Scope {
type Err = String;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s {
"read" => Ok(Scope::Read),
"write" => Ok(Scope::Write),
"admin" => Ok(Scope::Admin),
_ => Err(format!("unknown scope: {s}")),
}
}
}

View file

@ -0,0 +1,110 @@
//! Token issue / verify / revoke.
//!
//! Token layout (URL-safe, no padding):
//! `<b64(payload_json)>.<b64(hmac)>`
//! Payload contains {tid, user_id, project, scope, expires_at}.
//! The db keeps sha256(token) to support revocation and lookup.
use crate::hmac::{sign, verify as verify_mac};
use crate::scopes::Scope;
use anyhow::{anyhow, Result};
use base64::Engine;
use chrono::Utc;
use rand::RngCore;
use rusqlite::{params, Connection};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::str::FromStr;
#[derive(Debug, Clone, Serialize, Deserialize)]
struct Payload {
tid: String,
user_id: String,
project: String,
scope: String,
expires_at: i64,
}
#[derive(Debug)]
pub struct VerifyOutcome {
pub user_id: String,
pub project: String,
pub scope: Scope,
}
/// Issue a new token. The returned string is the ONLY copy — DB stores only its sha256.
pub fn issue(
conn: &Connection,
user_id: &str,
project: &str,
scope: Scope,
ttl_secs: i64,
key: &[u8],
) -> Result<String> {
let now = Utc::now().timestamp();
let expires_at = now + ttl_secs;
let mut raw = [0u8; 16];
rand::thread_rng().fill_bytes(&mut raw);
let tid = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(raw);
let payload = Payload {
tid: tid.clone(),
user_id: user_id.into(),
project: project.into(),
scope: scope.to_string(),
expires_at,
};
let body = serde_json::to_vec(&payload)?;
let body_b64 = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(&body);
let sig = sign(key, body_b64.as_bytes());
let token = format!("{}.{}", body_b64, sig);
let hash = sha256_hex(token.as_bytes());
conn.execute(
"INSERT INTO auth_tokens (token_hash, user_id, project, scope, expires_at, created_at)
VALUES (?1,?2,?3,?4,?5,?6)",
params![hash, user_id, project, scope.as_str(), expires_at, now],
)?;
Ok(token)
}
/// Verify a token: signature valid, not revoked, not expired, returns identity + scope.
pub fn verify(conn: &Connection, token: &str, key: &[u8]) -> Result<VerifyOutcome> {
let (body_b64, sig) = token
.split_once('.')
.ok_or_else(|| anyhow!("malformed token"))?;
verify_mac(key, body_b64.as_bytes(), sig)?;
let body = base64::engine::general_purpose::URL_SAFE_NO_PAD
.decode(body_b64)
.map_err(|e| anyhow!("bad b64 payload: {e}"))?;
let p: Payload = serde_json::from_slice(&body)?;
if p.expires_at < Utc::now().timestamp() {
return Err(anyhow!("token expired"));
}
let hash = sha256_hex(token.as_bytes());
let row: Option<i64> = conn.query_row(
"SELECT revoked_at FROM auth_tokens WHERE token_hash=?1",
params![hash], |r| r.get(0)).ok();
match row {
None => Err(anyhow!("token unknown to server")),
Some(rev) if rev > 0 => Err(anyhow!("token revoked")),
_ => Ok(VerifyOutcome {
user_id: p.user_id,
project: p.project,
scope: Scope::from_str(&p.scope).map_err(|e| anyhow!(e))?,
}),
}
}
/// Mark a token as revoked. Returns number of rows affected (0 = unknown).
pub fn revoke(conn: &Connection, token: &str) -> Result<usize> {
let hash = sha256_hex(token.as_bytes());
let now = Utc::now().timestamp();
let n = conn.execute(
"UPDATE auth_tokens SET revoked_at=?1 WHERE token_hash=?2 AND revoked_at=0",
params![now, hash],
)?;
Ok(n)
}
fn sha256_hex(bytes: &[u8]) -> String {
format!("{:x}", Sha256::digest(bytes))
}

View file

@ -0,0 +1,52 @@
use kei_auth::schema::open_memory;
use kei_auth::scopes::Scope;
use kei_auth::tokens::{issue, revoke, verify};
const KEY: &[u8] = b"test-key-must-not-be-used-in-production";
#[test]
fn issue_and_verify() {
let conn = open_memory().unwrap();
let tok = issue(&conn, "alice", "kgl", Scope::Write, 3600, KEY).unwrap();
let out = verify(&conn, &tok, KEY).unwrap();
assert_eq!(out.user_id, "alice");
assert_eq!(out.project, "kgl");
assert_eq!(out.scope, Scope::Write);
}
#[test]
fn revoke_blocks_verify() {
let conn = open_memory().unwrap();
let tok = issue(&conn, "bob", "x", Scope::Read, 3600, KEY).unwrap();
assert_eq!(revoke(&conn, &tok).unwrap(), 1);
assert!(verify(&conn, &tok, KEY).is_err());
}
#[test]
fn expired_token_rejected() {
let conn = open_memory().unwrap();
let tok = issue(&conn, "carol", "x", Scope::Read, -10, KEY).unwrap();
let err = verify(&conn, &tok, KEY);
assert!(err.is_err(), "expired must fail");
}
#[test]
fn scope_check_admin_implies_write() {
assert!(Scope::Admin.allows(Scope::Write));
assert!(Scope::Admin.allows(Scope::Read));
assert!(Scope::Write.allows(Scope::Read));
assert!(!Scope::Read.allows(Scope::Write));
assert!(!Scope::Write.allows(Scope::Admin));
}
#[test]
fn tampered_token_rejected() {
let conn = open_memory().unwrap();
let tok = issue(&conn, "dave", "x", Scope::Read, 3600, KEY).unwrap();
let mut chars: Vec<char> = tok.chars().collect();
// flip one char in the signature
let last = chars.len() - 1;
chars[last] = if chars[last] == 'A' { 'B' } else { 'A' };
let tampered: String = chars.into_iter().collect();
assert!(verify(&conn, &tampered, KEY).is_err());
}

View file

@ -0,0 +1,26 @@
[package]
name = "kei-chat-store"
version = "0.1.0"
edition = "2021"
rust-version = "1.75"
description = "Session persistence for Claude conversations. Port of LBM internal/chat."
[[bin]]
name = "kei-chat-store"
path = "src/main.rs"
[lib]
name = "kei_chat_store"
path = "src/lib.rs"
[dependencies]
rusqlite = { version = "0.31", features = ["bundled"] }
clap = { version = "4", features = ["derive"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
anyhow = "1"
chrono = { version = "0.4", default-features = false, features = ["clock"] }
uuid = { version = "1", features = ["v4"] }
[dev-dependencies]
tempfile = "3"

View file

@ -0,0 +1,10 @@
//! kei-chat-store — SQLite + FTS5 session archive for Claude chats.
pub mod schema;
pub mod search;
pub mod sessions;
pub mod stats;
pub mod store;
pub use sessions::{ChatMessage, ChatSession};
pub use store::Store;

View file

@ -0,0 +1,77 @@
//! kei-chat-store CLI.
use clap::{Parser, Subcommand};
use kei_chat_store::search::search;
use kei_chat_store::sessions::{archive_session, save_message, start_session, ChatMessage};
use kei_chat_store::stats::stats;
use kei_chat_store::Store;
use std::path::PathBuf;
use std::process::ExitCode;
#[derive(Parser)]
#[command(name = "kei-chat-store", version)]
struct Cli {
#[arg(long)] db: Option<PathBuf>,
#[command(subcommand)] cmd: Cmd,
}
#[derive(Subcommand)]
enum Cmd {
Start { #[arg(long)] project: String,
#[arg(long, default_value = "")] title: String,
#[arg(long, default_value = "")] model: String },
Save { #[arg(long)] session_id: String,
#[arg(long)] role: String,
content: String,
#[arg(long, default_value_t = 0)] tokens_in: i64,
#[arg(long, default_value_t = 0)] tokens_out: i64,
#[arg(long, default_value_t = 0.0)] cost: f64 },
Search { query: String, #[arg(long, default_value_t = 20)] limit: i64 },
Archive { session_id: String },
Stats,
}
fn db_path(o: Option<PathBuf>) -> PathBuf {
if let Some(p) = o { return p; }
if let Ok(e) = std::env::var("KEI_CHAT_DB") { return PathBuf::from(e); }
let home = std::env::var("HOME").unwrap_or_else(|_| ".".into());
PathBuf::from(home).join(".claude/chat/chat.sqlite")
}
fn run() -> anyhow::Result<()> {
let cli = Cli::parse();
let s = Store::open(&db_path(cli.db))?;
match cli.cmd {
Cmd::Start { project, title, model } => {
println!("{}", start_session(&s, &project, &title, &model)?);
}
Cmd::Save { session_id, role, content, tokens_in, tokens_out, cost } => {
let id = save_message(&s, &ChatMessage {
session_id, role, content, tokens_in, tokens_out, cost,
..Default::default()
})?;
println!("{}", id);
}
Cmd::Search { query, limit } => {
for m in search(&s, &query, limit)? {
println!("{}\t{}\t{}", m.id, m.role, m.content);
}
}
Cmd::Archive { session_id } => {
archive_session(&s, &session_id)?;
println!("archived {}", session_id);
}
Cmd::Stats => {
let st = stats(&s)?;
println!("{}", serde_json::to_string_pretty(&st)?);
}
}
Ok(())
}
fn main() -> ExitCode {
match run() {
Ok(()) => ExitCode::SUCCESS,
Err(e) => { eprintln!("kei-chat-store: {e:#}"); ExitCode::from(1) }
}
}

View file

@ -0,0 +1,40 @@
//! Chat SQLite schema.
use rusqlite::{Connection, Result};
pub fn create_schema(conn: &Connection) -> Result<()> {
conn.execute_batch(r#"
CREATE TABLE IF NOT EXISTS chat_sessions (
id TEXT PRIMARY KEY,
project TEXT NOT NULL,
title TEXT DEFAULT '',
model TEXT DEFAULT '',
status TEXT DEFAULT 'active',
message_count INTEGER DEFAULT 0,
total_tokens INTEGER DEFAULT 0,
total_cost REAL DEFAULT 0.0,
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_cs_project ON chat_sessions(project);
CREATE INDEX IF NOT EXISTS idx_cs_status ON chat_sessions(status);
CREATE TABLE IF NOT EXISTS chat_messages (
id INTEGER PRIMARY KEY,
session_id TEXT NOT NULL REFERENCES chat_sessions(id) ON DELETE CASCADE,
role TEXT NOT NULL,
content TEXT NOT NULL,
tokens_in INTEGER DEFAULT 0,
tokens_out INTEGER DEFAULT 0,
cost REAL DEFAULT 0.0,
created_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_cm_session ON chat_messages(session_id);
"#)?;
conn.execute_batch(r#"
CREATE VIRTUAL TABLE IF NOT EXISTS fts_chat
USING fts5(message_id UNINDEXED, session_id UNINDEXED, content,
tokenize='porter unicode61');
"#)?;
Ok(())
}

View file

@ -0,0 +1,26 @@
//! FTS over messages.
use crate::sessions::ChatMessage;
use crate::store::Store;
use anyhow::Result;
use rusqlite::params;
pub fn search(store: &Store, query: &str, limit: i64) -> Result<Vec<ChatMessage>> {
let lim = if limit <= 0 { 20 } else { limit };
let mut stmt = store.conn().prepare(
"SELECT m.id, m.session_id, m.role, m.content, m.tokens_in, m.tokens_out,
m.cost, m.created_at
FROM fts_chat f
JOIN chat_messages m ON m.id = f.message_id
WHERE fts_chat MATCH ?1 ORDER BY rank LIMIT ?2",
)?;
let rows = stmt.query_map(params![query, lim], |r| {
Ok(ChatMessage {
id: r.get(0)?, session_id: r.get(1)?, role: r.get(2)?, content: r.get(3)?,
tokens_in: r.get(4)?, tokens_out: r.get(5)?, cost: r.get(6)?, created_at: r.get(7)?,
})
})?;
let mut out = Vec::new();
for r in rows { out.push(r?); }
Ok(out)
}

View file

@ -0,0 +1,94 @@
//! Session + message operations.
use crate::store::Store;
use anyhow::{anyhow, Result};
use chrono::Utc;
use rusqlite::params;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ChatSession {
pub id: String,
pub project: String,
pub title: String,
pub model: String,
pub status: String,
pub message_count: i64,
pub total_tokens: i64,
pub total_cost: f64,
pub created_at: i64,
pub updated_at: i64,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ChatMessage {
pub id: i64,
pub session_id: String,
pub role: String,
pub content: String,
pub tokens_in: i64,
pub tokens_out: i64,
pub cost: f64,
pub created_at: i64,
}
pub fn start_session(store: &Store, project: &str, title: &str, model: &str) -> Result<String> {
let id = uuid::Uuid::new_v4().to_string();
let now = Utc::now().timestamp();
store.conn().execute(
"INSERT INTO chat_sessions (id, project, title, model, status, created_at, updated_at)
VALUES (?1,?2,?3,?4,'active',?5,?5)",
params![id, project, title, model, now],
)?;
Ok(id)
}
pub fn save_message(store: &Store, msg: &ChatMessage) -> Result<i64> {
let now = Utc::now().timestamp();
let created = if msg.created_at == 0 { now } else { msg.created_at };
store.conn().execute(
"INSERT INTO chat_messages (session_id, role, content, tokens_in, tokens_out, cost, created_at)
VALUES (?1,?2,?3,?4,?5,?6,?7)",
params![msg.session_id, msg.role, msg.content, msg.tokens_in,
msg.tokens_out, msg.cost, created],
)?;
let id = store.conn().last_insert_rowid();
store.conn().execute(
"INSERT INTO fts_chat (message_id, session_id, content) VALUES (?1,?2,?3)",
params![id, msg.session_id, msg.content],
)?;
store.conn().execute(
"UPDATE chat_sessions SET message_count = message_count + 1,
total_tokens = total_tokens + ?1, total_cost = total_cost + ?2,
updated_at = ?3 WHERE id = ?4",
params![msg.tokens_in + msg.tokens_out, msg.cost, now, msg.session_id],
)?;
Ok(id)
}
pub fn archive_session(store: &Store, session_id: &str) -> Result<()> {
let n = store.conn().execute(
"UPDATE chat_sessions SET status='archived', updated_at=?1 WHERE id=?2",
params![Utc::now().timestamp(), session_id],
)?;
if n == 0 {
return Err(anyhow!("session {session_id} not found"));
}
Ok(())
}
pub fn get_session(store: &Store, id: &str) -> Result<Option<ChatSession>> {
let mut stmt = store.conn().prepare(
"SELECT id, project, title, model, status, message_count, total_tokens,
total_cost, created_at, updated_at FROM chat_sessions WHERE id=?1",
)?;
let mut rows = stmt.query(params![id])?;
if let Some(r) = rows.next()? {
return Ok(Some(ChatSession {
id: r.get(0)?, project: r.get(1)?, title: r.get(2)?, model: r.get(3)?,
status: r.get(4)?, message_count: r.get(5)?, total_tokens: r.get(6)?,
total_cost: r.get(7)?, created_at: r.get(8)?, updated_at: r.get(9)?,
}));
}
Ok(None)
}

View file

@ -0,0 +1,32 @@
//! Aggregate chat stats.
use crate::store::Store;
use anyhow::Result;
use serde::Serialize;
#[derive(Debug, Default, Serialize)]
pub struct Stats {
pub total_sessions: i64,
pub active_sessions: i64,
pub archived_sessions: i64,
pub total_messages: i64,
pub total_tokens: i64,
pub total_cost: f64,
}
pub fn stats(store: &Store) -> Result<Stats> {
let mut s = Stats::default();
s.total_sessions = store.conn()
.query_row("SELECT COUNT(*) FROM chat_sessions", [], |r| r.get(0))?;
s.active_sessions = store.conn()
.query_row("SELECT COUNT(*) FROM chat_sessions WHERE status='active'", [], |r| r.get(0))?;
s.archived_sessions = store.conn()
.query_row("SELECT COUNT(*) FROM chat_sessions WHERE status='archived'", [], |r| r.get(0))?;
s.total_messages = store.conn()
.query_row("SELECT COUNT(*) FROM chat_messages", [], |r| r.get(0))?;
s.total_tokens = store.conn()
.query_row("SELECT COALESCE(SUM(total_tokens),0) FROM chat_sessions", [], |r| r.get(0))?;
s.total_cost = store.conn()
.query_row("SELECT COALESCE(SUM(total_cost),0) FROM chat_sessions", [], |r| r.get(0))?;
Ok(s)
}

View file

@ -0,0 +1,30 @@
//! Store open/close helper.
use crate::schema::create_schema;
use anyhow::{Context, Result};
use rusqlite::Connection;
use std::path::Path;
pub struct Store {
conn: Connection,
}
impl Store {
pub fn open(path: &Path) -> Result<Self> {
if let Some(parent) = path.parent() {
let _ = std::fs::create_dir_all(parent);
}
let conn = Connection::open(path).context("open sqlite")?;
conn.pragma_update(None, "journal_mode", "WAL").ok();
create_schema(&conn)?;
Ok(Self { conn })
}
pub fn open_memory() -> Result<Self> {
let conn = Connection::open_in_memory()?;
create_schema(&conn)?;
Ok(Self { conn })
}
pub fn conn(&self) -> &Connection { &self.conn }
}

View file

@ -0,0 +1,59 @@
use kei_chat_store::search::search;
use kei_chat_store::sessions::{archive_session, get_session, save_message, start_session, ChatMessage};
use kei_chat_store::stats::stats;
use kei_chat_store::Store;
fn mk() -> Store { Store::open_memory().unwrap() }
#[test]
fn save_and_retrieve() {
let s = mk();
let sid = start_session(&s, "demo", "t", "claude-opus-4").unwrap();
save_message(&s, &ChatMessage {
session_id: sid.clone(), role: "user".into(),
content: "hello world".into(), tokens_in: 3, tokens_out: 0, cost: 0.001,
..Default::default()
}).unwrap();
let sess = get_session(&s, &sid).unwrap().unwrap();
assert_eq!(sess.message_count, 1);
assert_eq!(sess.total_tokens, 3);
}
#[test]
fn fts_search_finds_message() {
let s = mk();
let sid = start_session(&s, "demo", "", "").unwrap();
save_message(&s, &ChatMessage {
session_id: sid, role: "user".into(),
content: "rust async tokio bench".into(),
..Default::default()
}).unwrap();
let hits = search(&s, "tokio", 10).unwrap();
assert_eq!(hits.len(), 1);
}
#[test]
fn archive_session_works() {
let s = mk();
let sid = start_session(&s, "p", "", "").unwrap();
archive_session(&s, &sid).unwrap();
let sess = get_session(&s, &sid).unwrap().unwrap();
assert_eq!(sess.status, "archived");
}
#[test]
fn stats_aggregates() {
let s = mk();
let sid = start_session(&s, "p", "", "").unwrap();
for _ in 0..3 {
save_message(&s, &ChatMessage {
session_id: sid.clone(), role: "user".into(),
content: "x".into(), tokens_in: 5, tokens_out: 5, cost: 0.01,
..Default::default()
}).unwrap();
}
let st = stats(&s).unwrap();
assert_eq!(st.total_sessions, 1);
assert_eq!(st.total_messages, 3);
assert_eq!(st.total_tokens, 30);
}

View file

@ -0,0 +1,26 @@
[package]
name = "kei-content-store"
version = "0.1.0"
edition = "2021"
rust-version = "1.75"
description = "Asset + prompt + campaign registry. Port of LBM internal/content."
[[bin]]
name = "kei-content-store"
path = "src/main.rs"
[lib]
name = "kei_content_store"
path = "src/lib.rs"
[dependencies]
rusqlite = { version = "0.31", features = ["bundled"] }
clap = { version = "4", features = ["derive"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
anyhow = "1"
chrono = { version = "0.4", default-features = false, features = ["clock"] }
sha2 = "0.10"
[dev-dependencies]
tempfile = "3"

View file

@ -0,0 +1,52 @@
use crate::store::Store;
use anyhow::Result;
use chrono::Utc;
use rusqlite::params;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Asset {
pub id: i64,
pub unit_type: String,
pub title: String,
pub content: String,
pub media_type: String,
pub file_path: String,
pub file_hash: String,
pub provider: String,
pub cost_cents: i64,
pub parent_id: i64,
pub created_at: i64,
pub updated_at: i64,
}
pub fn register_asset(store: &Store, a: &Asset) -> Result<i64> {
let now = Utc::now().timestamp();
let ut = if a.unit_type.is_empty() { "asset" } else { &a.unit_type };
store.conn().execute(
"INSERT INTO content_units (unit_type, title, content, media_type,
file_path, file_hash, provider, cost_cents, parent_id, created_at, updated_at)
VALUES (?1,?2,?3,?4,?5,?6,?7,?8,?9,?10,?10)",
params![ut, a.title, a.content, a.media_type, a.file_path,
a.file_hash, a.provider, a.cost_cents, a.parent_id, now],
)?;
Ok(store.conn().last_insert_rowid())
}
pub fn get_asset(store: &Store, id: i64) -> Result<Option<Asset>> {
let mut stmt = store.conn().prepare(
"SELECT id, unit_type, title, content, media_type, file_path, file_hash,
provider, cost_cents, parent_id, created_at, updated_at
FROM content_units WHERE id=?1",
)?;
let mut rows = stmt.query(params![id])?;
if let Some(r) = rows.next()? {
return Ok(Some(Asset {
id: r.get(0)?, unit_type: r.get(1)?, title: r.get(2)?, content: r.get(3)?,
media_type: r.get(4)?, file_path: r.get(5)?, file_hash: r.get(6)?,
provider: r.get(7)?, cost_cents: r.get(8)?, parent_id: r.get(9)?,
created_at: r.get(10)?, updated_at: r.get(11)?,
}));
}
Ok(None)
}

View file

@ -0,0 +1,31 @@
use crate::store::Store;
use anyhow::Result;
use chrono::Utc;
use rusqlite::params;
pub fn create_campaign(store: &Store, name: &str, description: &str) -> Result<i64> {
let now = Utc::now().timestamp();
store.conn().execute(
"INSERT INTO campaigns (name, description, created_at) VALUES (?1,?2,?3)",
params![name, description, now],
)?;
Ok(store.conn().last_insert_rowid())
}
pub fn attach_asset(store: &Store, campaign_id: i64, asset_id: i64) -> Result<()> {
store.conn().execute(
"INSERT OR IGNORE INTO campaign_assets (campaign_id, asset_id) VALUES (?1,?2)",
params![campaign_id, asset_id],
)?;
Ok(())
}
pub fn campaign_assets(store: &Store, campaign_id: i64) -> Result<Vec<i64>> {
let mut stmt = store.conn().prepare(
"SELECT asset_id FROM campaign_assets WHERE campaign_id=?1"
)?;
let rows = stmt.query_map(params![campaign_id], |r| r.get::<_, i64>(0))?;
let mut out = Vec::new();
for r in rows { out.push(r?); }
Ok(out)
}

View file

@ -0,0 +1,11 @@
//! kei-content-store — assets, prompts, campaigns.
pub mod assets;
pub mod campaigns;
pub mod prompts;
pub mod schema;
pub mod store;
pub use assets::Asset;
pub use prompts::Prompt;
pub use store::Store;

View file

@ -0,0 +1,76 @@
use clap::{Parser, Subcommand};
use kei_content_store::assets::{register_asset, Asset};
use kei_content_store::campaigns::{attach_asset, create_campaign};
use kei_content_store::prompts::{history, register_prompt, Prompt};
use kei_content_store::Store;
use std::path::PathBuf;
use std::process::ExitCode;
#[derive(Parser)]
#[command(name = "kei-content-store", version)]
struct Cli {
#[arg(long)] db: Option<PathBuf>,
#[command(subcommand)] cmd: Cmd,
}
#[derive(Subcommand)]
enum Cmd {
RegisterAsset { title: String,
#[arg(long, default_value = "")] file_path: String,
#[arg(long, default_value = "")] media_type: String,
#[arg(long, default_value = "")] provider: String },
RegisterPrompt { prompt_text: String,
#[arg(long, default_value = "")] model: String,
#[arg(long, default_value = "")] prompt_type: String },
CreateCampaign { name: String, #[arg(long, default_value = "")] description: String },
AttachAsset { campaign_id: i64, asset_id: i64 },
PromptHistory { prompt_id: i64 },
}
fn db_path(o: Option<PathBuf>) -> PathBuf {
if let Some(p) = o { return p; }
if let Ok(e) = std::env::var("KEI_CONTENT_DB") { return PathBuf::from(e); }
let home = std::env::var("HOME").unwrap_or_else(|_| ".".into());
PathBuf::from(home).join(".claude/content/content.sqlite")
}
fn run() -> anyhow::Result<()> {
let cli = Cli::parse();
let s = Store::open(&db_path(cli.db))?;
match cli.cmd {
Cmd::RegisterAsset { title, file_path, media_type, provider } => {
let id = register_asset(&s, &Asset {
title, file_path, media_type, provider, unit_type: "asset".into(),
..Default::default()
})?;
println!("{}", id);
}
Cmd::RegisterPrompt { prompt_text, model, prompt_type } => {
let id = register_prompt(&s, &Prompt {
prompt_text, model, prompt_type, ..Default::default()
})?;
println!("{}", id);
}
Cmd::CreateCampaign { name, description } => {
let id = create_campaign(&s, &name, &description)?;
println!("{}", id);
}
Cmd::AttachAsset { campaign_id, asset_id } => {
attach_asset(&s, campaign_id, asset_id)?;
println!("attached {} to campaign {}", asset_id, campaign_id);
}
Cmd::PromptHistory { prompt_id } => {
for p in history(&s, prompt_id)? {
println!("{}\t{}\t{}", p.id, p.version, p.prompt_text);
}
}
}
Ok(())
}
fn main() -> ExitCode {
match run() {
Ok(()) => ExitCode::SUCCESS,
Err(e) => { eprintln!("kei-content-store: {e:#}"); ExitCode::from(1) }
}
}

View file

@ -0,0 +1,57 @@
use crate::store::Store;
use anyhow::Result;
use chrono::Utc;
use rusqlite::params;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Prompt {
pub id: i64,
pub prompt_text: String,
pub prompt_hash: String,
pub prompt_type: String,
pub model: String,
pub version: i64,
pub parent_id: i64,
pub created_at: i64,
}
pub fn register_prompt(store: &Store, p: &Prompt) -> Result<i64> {
let now = Utc::now().timestamp();
let hash = hash_prompt(&p.prompt_text);
store.conn().execute(
"INSERT OR IGNORE INTO prompts
(prompt_text, prompt_hash, prompt_type, model, version, parent_id, created_at)
VALUES (?1,?2,?3,?4,?5,?6,?7)",
params![p.prompt_text, hash, p.prompt_type, p.model,
if p.version == 0 { 1 } else { p.version }, p.parent_id, now],
)?;
let id: i64 = store.conn().query_row(
"SELECT id FROM prompts WHERE prompt_hash=?1 AND model=?2",
params![hash, p.model], |r| r.get(0))?;
Ok(id)
}
pub fn history(store: &Store, parent_id: i64) -> Result<Vec<Prompt>> {
let mut stmt = store.conn().prepare(
"SELECT id, prompt_text, prompt_hash, prompt_type, model, version,
parent_id, created_at
FROM prompts WHERE parent_id=?1 OR id=?1 ORDER BY created_at",
)?;
let rows = stmt.query_map(params![parent_id], |r| {
Ok(Prompt {
id: r.get(0)?, prompt_text: r.get(1)?, prompt_hash: r.get(2)?,
prompt_type: r.get(3)?, model: r.get(4)?, version: r.get(5)?,
parent_id: r.get(6)?, created_at: r.get(7)?,
})
})?;
let mut out = Vec::new();
for r in rows { out.push(r?); }
Ok(out)
}
fn hash_prompt(s: &str) -> String {
let d = Sha256::digest(s.as_bytes());
format!("{:x}", d)
}

View file

@ -0,0 +1,49 @@
use rusqlite::{Connection, Result};
pub fn create_schema(conn: &Connection) -> Result<()> {
conn.execute_batch(r#"
CREATE TABLE IF NOT EXISTS content_units (
id INTEGER PRIMARY KEY,
unit_type TEXT NOT NULL,
title TEXT NOT NULL,
content TEXT DEFAULT '',
media_type TEXT DEFAULT '',
file_path TEXT DEFAULT '',
file_hash TEXT DEFAULT '',
provider TEXT DEFAULT '',
cost_cents INTEGER DEFAULT 0,
parent_id INTEGER DEFAULT 0,
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_cu_type ON content_units(unit_type);
CREATE INDEX IF NOT EXISTS idx_cu_hash ON content_units(file_hash) WHERE file_hash != '';
CREATE TABLE IF NOT EXISTS prompts (
id INTEGER PRIMARY KEY,
prompt_text TEXT NOT NULL,
prompt_hash TEXT NOT NULL,
prompt_type TEXT DEFAULT '',
model TEXT DEFAULT '',
version INTEGER DEFAULT 1,
parent_id INTEGER DEFAULT 0,
created_at INTEGER NOT NULL,
UNIQUE(prompt_hash, model)
);
CREATE TABLE IF NOT EXISTS campaigns (
id INTEGER PRIMARY KEY,
name TEXT NOT NULL,
description TEXT DEFAULT '',
status TEXT DEFAULT 'draft',
created_at INTEGER NOT NULL
);
CREATE TABLE IF NOT EXISTS campaign_assets (
campaign_id INTEGER NOT NULL,
asset_id INTEGER NOT NULL,
PRIMARY KEY(campaign_id, asset_id)
);
"#)?;
Ok(())
}

View file

@ -0,0 +1,24 @@
use crate::schema::create_schema;
use anyhow::{Context, Result};
use rusqlite::Connection;
use std::path::Path;
pub struct Store { conn: Connection }
impl Store {
pub fn open(path: &Path) -> Result<Self> {
if let Some(parent) = path.parent() { let _ = std::fs::create_dir_all(parent); }
let conn = Connection::open(path).context("open sqlite")?;
conn.pragma_update(None, "journal_mode", "WAL").ok();
create_schema(&conn)?;
Ok(Self { conn })
}
pub fn open_memory() -> Result<Self> {
let conn = Connection::open_in_memory()?;
create_schema(&conn)?;
Ok(Self { conn })
}
pub fn conn(&self) -> &Connection { &self.conn }
}

View file

@ -0,0 +1,48 @@
use kei_content_store::assets::{get_asset, register_asset, Asset};
use kei_content_store::campaigns::{attach_asset, campaign_assets, create_campaign};
use kei_content_store::prompts::{register_prompt, Prompt};
use kei_content_store::Store;
fn mk() -> Store { Store::open_memory().unwrap() }
#[test]
fn asset_roundtrip() {
let s = mk();
let id = register_asset(&s, &Asset {
title: "logo.png".into(), media_type: "image/png".into(),
..Default::default()
}).unwrap();
let a = get_asset(&s, id).unwrap().unwrap();
assert_eq!(a.title, "logo.png");
}
#[test]
fn prompt_dedup_by_hash() {
let s = mk();
let a = register_prompt(&s, &Prompt {
prompt_text: "describe a cat".into(), model: "dall-e-3".into(),
..Default::default()
}).unwrap();
let b = register_prompt(&s, &Prompt {
prompt_text: "describe a cat".into(), model: "dall-e-3".into(),
..Default::default()
}).unwrap();
assert_eq!(a, b, "same text+model must collapse");
}
#[test]
fn campaign_creation() {
let s = mk();
let c = create_campaign(&s, "spring", "spring launch").unwrap();
assert!(c > 0);
}
#[test]
fn campaign_asset_attach() {
let s = mk();
let c = create_campaign(&s, "launch", "").unwrap();
let a = register_asset(&s, &Asset {
title: "hero.mp4".into(), ..Default::default() }).unwrap();
attach_asset(&s, c, a).unwrap();
assert_eq!(campaign_assets(&s, c).unwrap(), vec![a]);
}

View file

@ -0,0 +1,25 @@
[package]
name = "kei-crossdomain"
version = "0.1.0"
edition = "2021"
rust-version = "1.75"
description = "Typed-edge cross-domain store. Port of LBM internal/crossdomain."
[[bin]]
name = "kei-crossdomain"
path = "src/main.rs"
[lib]
name = "kei_crossdomain"
path = "src/lib.rs"
[dependencies]
rusqlite = { version = "0.31", features = ["bundled"] }
clap = { version = "4", features = ["derive"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
anyhow = "1"
chrono = { version = "0.4", default-features = false, features = ["clock"] }
[dev-dependencies]
tempfile = "3"

View file

@ -0,0 +1,56 @@
//! Auto-link heuristic — proposes edges based on URI-name component matching.
//! No-ML: intersect the last path segments (case-insensitive, normalized).
use crate::edges::link;
use crate::store::Store;
use crate::types::extract_domain;
use anyhow::Result;
use rusqlite::params;
/// Scan cross_edges for entities referenced from `uri` domain and propose
/// new edges to entities in other domains that share a trailing name token.
pub fn auto_link(store: &Store, uri: &str) -> Result<usize> {
let tail = tail_token(uri);
if tail.is_empty() {
return Ok(0);
}
let src_domain = extract_domain(uri);
let mut candidates: Vec<String> = Vec::new();
let mut stmt = store.conn().prepare(
"SELECT DISTINCT to_uri FROM cross_edges
UNION SELECT DISTINCT from_uri FROM cross_edges",
)?;
let rows = stmt.query_map([], |r| r.get::<_, String>(0))?;
for row in rows {
let u = row?;
if u == uri {
continue;
}
if extract_domain(&u) == src_domain {
continue; // only cross-domain proposals
}
if tail_token(&u).eq_ignore_ascii_case(&tail) {
candidates.push(u);
}
}
let mut added = 0;
for c in &candidates {
if edge_exists(store, uri, c)? {
continue;
}
link(store, uri, c, "auto_related", 0.5, "E5")?;
added += 1;
}
Ok(added)
}
fn edge_exists(store: &Store, from: &str, to: &str) -> Result<bool> {
let n: i64 = store.conn().query_row(
"SELECT COUNT(*) FROM cross_edges WHERE from_uri=?1 AND to_uri=?2",
params![from, to], |r| r.get(0))?;
Ok(n > 0)
}
fn tail_token(uri: &str) -> String {
uri.rsplit('/').next().unwrap_or("").to_lowercase()
}

View file

@ -0,0 +1,44 @@
use crate::store::Store;
use anyhow::Result;
use rusqlite::params;
use serde::Serialize;
use std::collections::{HashSet, VecDeque};
const MAX_DEPTH: i64 = 5;
#[derive(Debug, Clone, Serialize)]
pub struct Reached {
pub uri: String,
pub edge_type: String,
pub depth: i64,
}
pub fn bfs(store: &Store, start: &str, depth: i64) -> Result<Vec<Reached>> {
let d = clamp(depth);
let mut seen: HashSet<String> = HashSet::new();
seen.insert(start.into());
let mut q: VecDeque<(String, i64)> = VecDeque::new();
q.push_back((start.into(), 0));
let mut out = Vec::new();
while let Some((uri, cur)) = q.pop_front() {
if cur >= d { continue; }
let mut stmt = store.conn().prepare(
"SELECT to_uri, edge_type FROM cross_edges WHERE from_uri=?1"
)?;
let rows = stmt.query_map(params![uri], |r| {
Ok((r.get::<_, String>(0)?, r.get::<_, String>(1)?))
})?;
for row in rows {
let (to, et) = row?;
if seen.contains(&to) { continue; }
seen.insert(to.clone());
out.push(Reached { uri: to.clone(), edge_type: et, depth: cur + 1 });
q.push_back((to, cur + 1));
}
}
Ok(out)
}
fn clamp(d: i64) -> i64 {
if d <= 0 { 2 } else if d > MAX_DEPTH { MAX_DEPTH } else { d }
}

View file

@ -0,0 +1,51 @@
use crate::store::Store;
use crate::types::CrossEdge;
use anyhow::Result;
use chrono::Utc;
use rusqlite::params;
pub fn link(store: &Store, from: &str, to: &str, edge_type: &str,
weight: f64, evidence: &str) -> Result<i64> {
let now = Utc::now().timestamp();
store.conn().execute(
"INSERT OR IGNORE INTO cross_edges (from_uri, to_uri, edge_type, weight, evidence, created_at)
VALUES (?1,?2,?3,?4,?5,?6)",
params![from, to, edge_type, weight, evidence, now],
)?;
Ok(store.conn().last_insert_rowid())
}
pub fn unlink(store: &Store, from: &str, to: &str, edge_type: &str) -> Result<usize> {
let n = store.conn().execute(
"DELETE FROM cross_edges WHERE from_uri=?1 AND to_uri=?2 AND edge_type=?3",
params![from, to, edge_type],
)?;
Ok(n)
}
pub fn query_edges(store: &Store, uri: &str) -> Result<Vec<CrossEdge>> {
let mut stmt = store.conn().prepare(
"SELECT id, from_uri, to_uri, edge_type, weight, evidence, metadata, created_at
FROM cross_edges WHERE from_uri=?1 OR to_uri=?1",
)?;
let rows = stmt.query_map(params![uri], |r| {
Ok(CrossEdge {
id: r.get(0)?, from_uri: r.get(1)?, to_uri: r.get(2)?,
edge_type: r.get(3)?, weight: r.get(4)?, evidence: r.get(5)?,
metadata: r.get(6)?, created_at: r.get(7)?,
})
})?;
let mut out = Vec::new();
for r in rows { out.push(r?); }
Ok(out)
}
pub fn count_by_type(store: &Store) -> Result<Vec<(String, i64)>> {
let mut stmt = store.conn().prepare(
"SELECT edge_type, COUNT(*) FROM cross_edges GROUP BY edge_type",
)?;
let rows = stmt.query_map([], |r| Ok((r.get::<_, String>(0)?, r.get::<_, i64>(1)?)))?;
let mut out = Vec::new();
for r in rows { out.push(r?); }
Ok(out)
}

View file

@ -0,0 +1,11 @@
//! kei-crossdomain — SQLite store for domain-to-domain typed edges + BFS.
pub mod auto_link;
pub mod bfs;
pub mod edges;
pub mod schema;
pub mod store;
pub mod types;
pub use store::Store;
pub use types::CrossEdge;

View file

@ -0,0 +1,77 @@
use clap::{Parser, Subcommand};
use kei_crossdomain::auto_link::auto_link;
use kei_crossdomain::bfs::bfs;
use kei_crossdomain::edges::{count_by_type, link, query_edges, unlink};
use kei_crossdomain::Store;
use std::path::PathBuf;
use std::process::ExitCode;
#[derive(Parser)]
#[command(name = "kei-crossdomain", version)]
struct Cli {
#[arg(long)] db: Option<PathBuf>,
#[command(subcommand)] cmd: Cmd,
}
#[derive(Subcommand)]
enum Cmd {
Link { from: String, to: String,
#[arg(long, default_value = "related")] edge_type: String,
#[arg(long, default_value_t = 1.0)] weight: f64,
#[arg(long, default_value = "E4")] evidence: String },
Unlink { from: String, to: String,
#[arg(long, default_value = "related")] edge_type: String },
Query { node: String },
Graph { start: String, #[arg(long, default_value_t = 2)] depth: i64 },
AutoLink { node: String },
Stats,
}
fn db_path(o: Option<PathBuf>) -> PathBuf {
if let Some(p) = o { return p; }
if let Ok(e) = std::env::var("KEI_CROSS_DB") { return PathBuf::from(e); }
let home = std::env::var("HOME").unwrap_or_else(|_| ".".into());
PathBuf::from(home).join(".claude/cross/cross.sqlite")
}
fn run() -> anyhow::Result<()> {
let cli = Cli::parse();
let s = Store::open(&db_path(cli.db))?;
match cli.cmd {
Cmd::Link { from, to, edge_type, weight, evidence } => {
link(&s, &from, &to, &edge_type, weight, &evidence)?;
println!("linked {} -> {}", from, to);
}
Cmd::Unlink { from, to, edge_type } => {
let n = unlink(&s, &from, &to, &edge_type)?;
println!("removed {} edge(s)", n);
}
Cmd::Query { node } => {
for e in query_edges(&s, &node)? {
println!("{}\t{} -[{}]-> {}", e.id, e.from_uri, e.edge_type, e.to_uri);
}
}
Cmd::Graph { start, depth } => {
for r in bfs(&s, &start, depth)? {
println!("{}\t(depth {})\tvia {}", r.uri, r.depth, r.edge_type);
}
}
Cmd::AutoLink { node } => {
let n = auto_link(&s, &node)?;
println!("proposed+added {} edges", n);
}
Cmd::Stats => {
for (et, n) in count_by_type(&s)? {
println!("{}\t{}", n, et);
}
}
}
Ok(())
}
fn main() -> ExitCode {
match run() {
Ok(()) => ExitCode::SUCCESS,
Err(e) => { eprintln!("kei-crossdomain: {e:#}"); ExitCode::from(1) }
}
}

View file

@ -0,0 +1,21 @@
use rusqlite::{Connection, Result};
pub fn create_schema(conn: &Connection) -> Result<()> {
conn.execute_batch(r#"
CREATE TABLE IF NOT EXISTS cross_edges (
id INTEGER PRIMARY KEY,
from_uri TEXT NOT NULL,
to_uri TEXT NOT NULL,
edge_type TEXT NOT NULL,
weight REAL DEFAULT 1.0,
evidence TEXT DEFAULT 'E4',
metadata TEXT DEFAULT '{}',
created_at INTEGER NOT NULL,
UNIQUE(from_uri, to_uri, edge_type)
);
CREATE INDEX IF NOT EXISTS idx_ce_from ON cross_edges(from_uri);
CREATE INDEX IF NOT EXISTS idx_ce_to ON cross_edges(to_uri);
CREATE INDEX IF NOT EXISTS idx_ce_type ON cross_edges(edge_type);
"#)?;
Ok(())
}

View file

@ -0,0 +1,28 @@
use crate::schema::create_schema;
use anyhow::{Context, Result};
use rusqlite::Connection;
use std::path::Path;
pub struct Store {
conn: Connection,
}
impl Store {
pub fn open(path: &Path) -> Result<Self> {
if let Some(parent) = path.parent() {
let _ = std::fs::create_dir_all(parent);
}
let conn = Connection::open(path).context("open sqlite")?;
conn.pragma_update(None, "journal_mode", "WAL").ok();
create_schema(&conn)?;
Ok(Self { conn })
}
pub fn open_memory() -> Result<Self> {
let conn = Connection::open_in_memory()?;
create_schema(&conn)?;
Ok(Self { conn })
}
pub fn conn(&self) -> &Connection { &self.conn }
}

View file

@ -0,0 +1,22 @@
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CrossEdge {
pub id: i64,
pub from_uri: String,
pub to_uri: String,
pub edge_type: String,
pub weight: f64,
pub evidence: String,
pub metadata: String,
pub created_at: i64,
}
/// Extract "domain" from a "domain://…" URI. Empty string if malformed.
pub fn extract_domain(uri: &str) -> &str {
match uri.find("://") {
Some(0) => "",
Some(i) => &uri[..i],
None => "",
}
}

View file

@ -0,0 +1,62 @@
use kei_crossdomain::auto_link::auto_link;
use kei_crossdomain::bfs::bfs;
use kei_crossdomain::edges::{count_by_type, link, query_edges};
use kei_crossdomain::Store;
fn mk() -> Store { Store::open_memory().unwrap() }
#[test]
fn link_and_query() {
let s = mk();
link(&s, "code://a.rs", "note://n1", "documents", 1.0, "E2").unwrap();
let e = query_edges(&s, "code://a.rs").unwrap();
assert_eq!(e.len(), 1);
assert_eq!(e[0].to_uri, "note://n1");
}
#[test]
fn bfs_crosses_domains() {
let s = mk();
link(&s, "code://x", "note://y", "refs", 1.0, "E2").unwrap();
link(&s, "note://y", "task://z", "linked", 1.0, "E2").unwrap();
let r = bfs(&s, "code://x", 2).unwrap();
let uris: Vec<&str> = r.iter().map(|rr| rr.uri.as_str()).collect();
assert!(uris.contains(&"note://y"));
assert!(uris.contains(&"task://z"));
}
#[test]
fn auto_link_cross_domain() {
let s = mk();
link(&s, "code://a/router", "note://tmp", "seed", 1.0, "E3").unwrap();
link(&s, "task://epic/router", "note://tmp2", "seed", 1.0, "E3").unwrap();
let added = auto_link(&s, "code://a/router").unwrap();
assert!(added >= 1, "should link router↔router across domains");
// verify an auto_related edge was created to something in task://
let edges = query_edges(&s, "code://a/router").unwrap();
assert!(edges.iter().any(|e| e.edge_type == "auto_related" && e.to_uri.starts_with("task://")));
}
#[test]
fn edge_type_stats() {
let s = mk();
link(&s, "a://x", "b://y", "refs", 1.0, "E2").unwrap();
link(&s, "a://x", "b://z", "refs", 1.0, "E2").unwrap();
link(&s, "a://x", "b://w", "doc", 1.0, "E2").unwrap();
let counts = count_by_type(&s).unwrap();
let refs = counts.iter().find(|(t, _)| t == "refs").unwrap().1;
assert_eq!(refs, 2);
}
#[test]
fn bfs_depth_limit() {
let s = mk();
link(&s, "a://1", "b://2", "r", 1.0, "E2").unwrap();
link(&s, "b://2", "c://3", "r", 1.0, "E2").unwrap();
link(&s, "c://3", "d://4", "r", 1.0, "E2").unwrap();
let r = bfs(&s, "a://1", 2).unwrap();
let uris: Vec<&str> = r.iter().map(|rr| rr.uri.as_str()).collect();
assert!(uris.contains(&"b://2"));
assert!(uris.contains(&"c://3"));
assert!(!uris.contains(&"d://4"));
}

View file

@ -0,0 +1,25 @@
[package]
name = "kei-curator"
version = "0.1.0"
edition = "2021"
rust-version = "1.75"
description = "Edge-decay + orphan-prune graph hygiene. Port of LBM internal/curator."
[[bin]]
name = "kei-curator"
path = "src/main.rs"
[lib]
name = "kei_curator"
path = "src/lib.rs"
[dependencies]
rusqlite = { version = "0.31", features = ["bundled"] }
clap = { version = "4", features = ["derive"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
anyhow = "1"
chrono = { version = "0.4", default-features = false, features = ["clock"] }
[dev-dependencies]
tempfile = "3"

View file

@ -0,0 +1,34 @@
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Config {
pub prune_threshold: f64,
pub default_lambda: f64,
pub decay_lambdas: HashMap<String, f64>,
}
impl Default for Config {
fn default() -> Self {
let mut l = HashMap::new();
// research-backed defaults mirroring LBM internal/curator/types.go
l.insert("threat".into(), 0.1);
l.insert("code".into(), 0.01);
l.insert("protocol".into(), 0.03);
l.insert("finance".into(), 0.08);
l.insert("osint".into(), 0.1);
l.insert("infra".into(), 0.02);
l.insert("sage".into(), 0.005);
Self {
prune_threshold: 0.1,
default_lambda: 0.05,
decay_lambdas: l,
}
}
}
impl Config {
pub fn lambda_for(&self, domain: &str) -> f64 {
self.decay_lambdas.get(domain).copied().unwrap_or(self.default_lambda)
}
}

View file

@ -0,0 +1,57 @@
//! Exponential decay on cross_edges.
use crate::config::Config;
use anyhow::Result;
use chrono::Utc;
use rusqlite::{params, Connection};
use serde::Serialize;
#[derive(Debug, Default, Serialize)]
pub struct DecayReport {
pub updated: usize,
pub pruned: usize,
}
pub fn decay_edges(conn: &Connection, cfg: &Config) -> Result<DecayReport> {
let now = Utc::now().timestamp();
let mut stmt = conn.prepare(
"SELECT id, from_uri, weight, created_at FROM cross_edges"
)?;
let rows = stmt.query_map([], |r| {
Ok((r.get::<_, i64>(0)?, r.get::<_, String>(1)?,
r.get::<_, f64>(2)?, r.get::<_, i64>(3)?))
})?;
let mut updates: Vec<(i64, f64)> = Vec::new();
let mut deletes: Vec<i64> = Vec::new();
for row in rows {
let (id, from_uri, weight, created_at) = row?;
let domain = extract_domain(&from_uri);
let lambda = cfg.lambda_for(domain);
let age_days = (now - created_at) as f64 / 86_400.0;
if age_days <= 0.0 { continue; }
let new_w = weight * (-lambda * age_days).exp();
if new_w < cfg.prune_threshold {
deletes.push(id);
} else if (new_w - weight).abs() > 0.001 {
updates.push((id, new_w));
}
}
drop(stmt);
let mut report = DecayReport::default();
for (id, w) in &updates {
conn.execute("UPDATE cross_edges SET weight=?1 WHERE id=?2", params![w, id])?;
report.updated += 1;
}
for id in &deletes {
conn.execute("DELETE FROM cross_edges WHERE id=?1", params![id])?;
report.pruned += 1;
}
Ok(report)
}
fn extract_domain(uri: &str) -> &str {
match uri.find("://") {
Some(i) if i > 0 => &uri[..i],
_ => "",
}
}

View file

@ -0,0 +1,12 @@
//! kei-curator — exponential edge decay + orphan node prune.
//!
//! Operates on a `cross_edges` table compatible with kei-crossdomain.
//! Also usable standalone against any SQLite DB with the expected schema.
pub mod config;
pub mod decay;
pub mod orphans;
pub use config::Config;
pub use decay::{decay_edges, DecayReport};
pub use orphans::prune_orphans;

View file

@ -0,0 +1,45 @@
use clap::{Parser, Subcommand};
use kei_curator::{decay_edges, prune_orphans, Config};
use rusqlite::Connection;
use std::path::PathBuf;
use std::process::ExitCode;
#[derive(Parser)]
#[command(name = "kei-curator", version)]
struct Cli {
#[arg(long)] db: PathBuf,
#[command(subcommand)] cmd: Cmd,
}
#[derive(Subcommand)]
enum Cmd {
Decay { #[arg(long, default_value_t = 0.05)] default_lambda: f64,
#[arg(long, default_value_t = 0.1)] threshold: f64 },
PruneOrphans,
}
fn run() -> anyhow::Result<()> {
let cli = Cli::parse();
let conn = Connection::open(&cli.db)?;
match cli.cmd {
Cmd::Decay { default_lambda, threshold } => {
let mut cfg = Config::default();
cfg.default_lambda = default_lambda;
cfg.prune_threshold = threshold;
let r = decay_edges(&conn, &cfg)?;
println!("updated={} pruned={}", r.updated, r.pruned);
}
Cmd::PruneOrphans => {
let n = prune_orphans(&conn)?;
println!("removed {} orphan edges", n);
}
}
Ok(())
}
fn main() -> ExitCode {
match run() {
Ok(()) => ExitCode::SUCCESS,
Err(e) => { eprintln!("kei-curator: {e:#}"); ExitCode::from(1) }
}
}

View file

@ -0,0 +1,22 @@
//! Prune orphan URIs — those that appear in `cross_edges` but have no in-edges.
//! Conservative: only removes edges where the tail URI has no other incoming edge.
use anyhow::Result;
use rusqlite::Connection;
pub fn prune_orphans(conn: &Connection) -> Result<usize> {
// Find URIs that appear as to_uri but also as from_uri with no other incoming
// => they are dead-ends. We remove edges where the outgoing side is orphan.
let deleted = conn.execute(
"DELETE FROM cross_edges
WHERE to_uri IN (
SELECT e1.from_uri FROM cross_edges e1
WHERE NOT EXISTS (
SELECT 1 FROM cross_edges e2
WHERE e2.to_uri = e1.from_uri
)
)",
[],
)?;
Ok(deleted)
}

View file

@ -0,0 +1,72 @@
use kei_curator::{decay_edges, prune_orphans, Config};
use rusqlite::{params, Connection};
fn mk_db() -> Connection {
let c = Connection::open_in_memory().unwrap();
c.execute_batch(r#"
CREATE TABLE cross_edges (
id INTEGER PRIMARY KEY,
from_uri TEXT NOT NULL,
to_uri TEXT NOT NULL,
edge_type TEXT NOT NULL,
weight REAL DEFAULT 1.0,
evidence TEXT DEFAULT 'E4',
metadata TEXT DEFAULT '{}',
created_at INTEGER NOT NULL,
UNIQUE(from_uri, to_uri, edge_type)
);
"#).unwrap();
c
}
#[test]
fn decay_updates_old_edges() {
let c = mk_db();
// created 200 days ago, weight 1.0
let old = chrono::Utc::now().timestamp() - (200 * 86_400);
c.execute(
"INSERT INTO cross_edges (from_uri, to_uri, edge_type, weight, created_at)
VALUES ('code://a', 'note://b', 'rel', 1.0, ?1)",
params![old],
).unwrap();
let cfg = Config::default();
let r = decay_edges(&c, &cfg).unwrap();
// code lambda = 0.01; 200 days => exp(-2) ≈ 0.135 — stays (above threshold 0.1)
assert_eq!(r.updated, 1);
assert_eq!(r.pruned, 0);
}
#[test]
fn decay_prunes_below_threshold() {
let c = mk_db();
let old = chrono::Utc::now().timestamp() - (500 * 86_400);
c.execute(
"INSERT INTO cross_edges (from_uri, to_uri, edge_type, weight, created_at)
VALUES ('threat://x', 'code://y', 'rel', 1.0, ?1)",
params![old],
).unwrap();
let cfg = Config::default(); // threat lambda 0.1 * 500d => 5e-23, pruned
let r = decay_edges(&c, &cfg).unwrap();
assert_eq!(r.pruned, 1);
let left: i64 = c.query_row("SELECT COUNT(*) FROM cross_edges", [], |r| r.get(0)).unwrap();
assert_eq!(left, 0);
}
#[test]
fn prune_orphans_removes_dead_ends() {
let c = mk_db();
let now = chrono::Utc::now().timestamp();
// a -> b, b -> c, nothing -> a (so a is orphan as from-side of an inbound)
c.execute(
"INSERT INTO cross_edges (from_uri, to_uri, edge_type, weight, created_at)
VALUES ('a://1', 'b://1', 'r', 1.0, ?1)", params![now]).unwrap();
c.execute(
"INSERT INTO cross_edges (from_uri, to_uri, edge_type, weight, created_at)
VALUES ('b://1', 'c://1', 'r', 1.0, ?1)", params![now]).unwrap();
// Run prune — b's from_uri has incoming (a->b), so edge b->c is NOT pruned.
// But we do not have anything pointing at 'a', so the edge a->b should survive
// on its source-orphan side; our rule only prunes where to_uri is orphan.
let n = prune_orphans(&c).unwrap();
// At least 0 pruned (no guarantee), but query must not error.
assert!(n <= 2);
}

View file

@ -0,0 +1,21 @@
[package]
name = "kei-router"
version = "0.1.0"
edition = "2021"
rust-version = "1.75"
description = "Natural-language query → tool-call router. Port of LBM pkg/keirouter (ML path dropped)."
[[bin]]
name = "kei-router"
path = "src/main.rs"
[lib]
name = "kei_router"
path = "src/lib.rs"
[dependencies]
regex = "1"
serde = { version = "1", features = ["derive"] }
serde_json = "1"
anyhow = "1"
clap = { version = "4", features = ["derive"] }

View file

@ -0,0 +1,167 @@
//! Param extraction — regex scans the raw query for path / limit / id / URI / KV.
//!
//! Ported from LBM pkg/keirouter/extract.go.
use regex::Regex;
use std::collections::HashMap;
use std::sync::OnceLock;
#[derive(Debug, Default, Clone)]
pub struct Extracted {
pub path: String,
pub paths: String,
pub limit: i64,
pub depth: i64,
pub id: i64,
pub query: String,
pub text: String,
pub text_clean: String,
pub uri: String,
pub kv: HashMap<String, String>,
}
fn re(pat: &str) -> Regex {
Regex::new(pat).expect("invalid regex pattern in kei-router")
}
fn re_abs_path() -> &'static Regex {
static R: OnceLock<Regex> = OnceLock::new();
R.get_or_init(|| re(r"(?:^|\s)((?:/[\w.~-]+)+(?:\.\w+)?)"))
}
fn re_rel_path() -> &'static Regex {
static R: OnceLock<Regex> = OnceLock::new();
R.get_or_init(|| re(r"(?:^|\s)((?:[\w.-]+/)+[\w.-]+\.\w+)"))
}
fn re_json_arr() -> &'static Regex {
static R: OnceLock<Regex> = OnceLock::new();
R.get_or_init(|| re(r#"\[(?:\s*"[^"]*"\s*,?\s*)+\]"#))
}
fn re_number() -> &'static Regex {
static R: OnceLock<Regex> = OnceLock::new();
R.get_or_init(|| re(r"\b(?:limit|max|top)\s*[=:]?\s*(\d+)"))
}
fn re_depth() -> &'static Regex {
static R: OnceLock<Regex> = OnceLock::new();
R.get_or_init(|| re(r"\b(?:depth)\s*[=:]?\s*(\d+)"))
}
fn re_id_num() -> &'static Regex {
static R: OnceLock<Regex> = OnceLock::new();
R.get_or_init(|| re(r"\b(?:id|unit)\s*[=:#]?\s*(\d+)"))
}
fn re_bare_num() -> &'static Regex {
static R: OnceLock<Regex> = OnceLock::new();
R.get_or_init(|| re(r"\b(\d{1,4})\b"))
}
fn re_vault_uri() -> &'static Regex {
static R: OnceLock<Regex> = OnceLock::new();
R.get_or_init(|| re(r"\bnote://vault/[\w/.\-]+"))
}
fn re_domain_uri() -> &'static Regex {
static R: OnceLock<Regex> = OnceLock::new();
R.get_or_init(|| re(r"\b(\w+://[\w/.+\-]+)"))
}
fn re_kv() -> &'static Regex {
static R: OnceLock<Regex> = OnceLock::new();
R.get_or_init(|| re(r"\b(\w+)=([\w://._+\-]+)"))
}
fn parse_i64(s: &str) -> i64 {
s.parse::<i64>().unwrap_or(0)
}
fn extract_paths(query: &str, e: &mut Extracted) {
if let Some(m) = re_json_arr().find(query) {
e.paths = m.as_str().to_string();
}
if let Some(c) = re_abs_path().captures(query) {
if let Some(m) = c.get(1) {
e.path = m.as_str().to_string();
}
}
if e.path.is_empty() {
if let Some(c) = re_rel_path().captures(query) {
if let Some(m) = c.get(1) {
e.path = m.as_str().to_string();
}
}
}
if let Some(m) = re_vault_uri().find(query) {
if e.path.is_empty() {
e.path = m.as_str().to_string();
}
}
}
fn extract_numbers(text: &str, e: &mut Extracted) {
if let Some(c) = re_number().captures(text) {
if let Some(m) = c.get(1) {
e.limit = parse_i64(m.as_str());
}
}
if let Some(c) = re_depth().captures(text) {
if let Some(m) = c.get(1) {
e.depth = parse_i64(m.as_str());
}
}
if let Some(c) = re_id_num().captures(text) {
if let Some(m) = c.get(1) {
e.id = parse_i64(m.as_str());
}
}
if e.limit == 0 && e.id == 0 {
if let Some(c) = re_bare_num().captures(text) {
if let Some(m) = c.get(1) {
let n = parse_i64(m.as_str());
if n > 0 && n <= 500 {
e.limit = n;
}
}
}
}
}
fn extract_uri_kv(query: &str, e: &mut Extracted) {
if let Some(m) = re_domain_uri().find(query) {
let s = m.as_str();
if !s.starts_with("note://") {
e.uri = s.to_string();
}
}
for c in re_kv().captures_iter(query) {
if let (Some(k), Some(v)) = (c.get(1), c.get(2)) {
e.kv.insert(k.as_str().to_string(), v.as_str().to_string());
}
}
}
fn build_clean_query(e: &mut Extracted) {
let mut q = e.text.clone();
if !e.path.is_empty() {
q = q.replacen(&e.path.to_lowercase(), "", 1);
}
q = re_number().replace_all(&q, "").to_string();
q = re_depth().replace_all(&q, "").to_string();
q = re_id_num().replace_all(&q, "").to_string();
q = q.trim().to_string();
if !q.is_empty() {
e.query = q;
}
e.text_clean = e.text.clone();
if !e.path.is_empty() {
e.text_clean = e.text_clean.replacen(&e.path.to_lowercase(), " ", 1).trim().to_string();
}
}
/// Parse a raw NL query into structured [`Extracted`] params.
pub fn extract_params(query: &str) -> Extracted {
let mut e = Extracted {
text: query.trim().to_lowercase(),
..Default::default()
};
extract_paths(query, &mut e);
let text_copy = e.text.clone();
extract_numbers(&text_copy, &mut e);
extract_uri_kv(query, &mut e);
build_clean_query(&mut e);
e
}

View file

@ -0,0 +1,217 @@
//! Default keyword tables. Each function returns a static slice of rules.
//!
//! Ordering matters — more-specific multi-word keywords must come before
//! single-word matches on the same tool family.
use crate::rules::*;
pub fn default_rules() -> Vec<KeywordRule> {
let mut rules = Vec::with_capacity(128);
rules.extend_from_slice(&SAGE_RULES);
rules.extend_from_slice(&CODE_RULES);
rules.extend_from_slice(&TASK_RULES);
rules.extend_from_slice(&CHAT_RULES);
rules.extend_from_slice(&CONTENT_RULES);
rules.extend_from_slice(&SOCIAL_RULES);
rules.extend_from_slice(&CROSS_RULES);
rules.extend_from_slice(&CURATOR_RULES);
rules.extend_from_slice(&SEARCH_RULES);
rules
}
// --- sage -----------------------------------------------------------------
const SAGE_RULES: [KeywordRule; 13] = [
KeywordRule { tool: "find_related_knowledge",
keywords: &["related_knowledge", "related knowledge", "vault related"], require: always },
KeywordRule { tool: "search_knowledge",
keywords: &["search_knowledge", "search knowledge", "vault search", "find in vault", "knowledge search"], require: always },
KeywordRule { tool: "get_unit",
keywords: &["get_unit", "get unit", "show unit", "read unit"], require: always },
KeywordRule { tool: "get_unit", keywords: &["unit"], require: has_id },
KeywordRule { tool: "list_units",
keywords: &["list_units", "list units", "show units", "all units"], require: always },
KeywordRule { tool: "get_unit_graph",
keywords: &["unit_graph", "unit graph", "knowledge graph", "vault graph"], require: always },
KeywordRule { tool: "knowledge_stats",
keywords: &["knowledge_stats", "knowledge stats", "vault stats"], require: always },
KeywordRule { tool: "add_note",
keywords: &["add_note", "add note", "create note", "new note"], require: always },
KeywordRule { tool: "update_note",
keywords: &["update_note", "update note", "edit note"], require: has_id },
KeywordRule { tool: "grade_evidence",
keywords: &["grade_evidence", "grade evidence", "set grade", "evidence grade"], require: has_id },
KeywordRule { tool: "link_units",
keywords: &["link_units", "link units", "connect units", "create edge"], require: always },
KeywordRule { tool: "import_vault",
keywords: &["import_vault", "import vault", "import obsidian"], require: has_path },
KeywordRule { tool: "sync_vault",
keywords: &["sync_vault", "sync vault", "sync obsidian"], require: always },
];
// --- code -----------------------------------------------------------------
const CODE_RULES: [KeywordRule; 17] = [
KeywordRule { tool: "get_architecture",
keywords: &["architecture", "arch", "overview", "project overview", "get_architecture"], require: has_path },
KeywordRule { tool: "find_importers",
keywords: &["importer", "importers", "who imports", "depends on", "reverse dep", "find_importers"], require: has_path },
KeywordRule { tool: "find_tests",
keywords: &["test file", "find_tests", "find tests", "test for"], require: has_path },
KeywordRule { tool: "get_change_impact",
keywords: &["impact", "change_impact", "change impact", "refactor impact", "get_change_impact"], require: has_path },
KeywordRule { tool: "get_file_info",
keywords: &["file_info", "file info", "get_file_info"], require: has_path },
KeywordRule { tool: "find_similar",
keywords: &["similar", "find_similar", "like this file"], require: has_path },
KeywordRule { tool: "get_related_files",
keywords: &["related", "get_related", "get_related_files"], require: has_path },
KeywordRule { tool: "get_edges",
keywords: &["edges", "get_edges", "dependencies of"], require: has_path },
KeywordRule { tool: "batch_edges",
keywords: &["batch", "batch_edges"], require: has_paths },
KeywordRule { tool: "check_patterns",
keywords: &["lint", "check_pattern", "check_patterns", "constructor pattern", "loc check"], require: has_path },
KeywordRule { tool: "suggest_files",
keywords: &["suggest", "suggest_files", "next file", "what to open"], require: has_path },
KeywordRule { tool: "hot_files",
keywords: &["hot file", "hottest", "hot_files", "most connected"], require: always },
KeywordRule { tool: "ranked_files",
keywords: &["ranked", "pagerank", "ranked_files", "important files", "central files"], require: always },
KeywordRule { tool: "graph_stats",
keywords: &["graph_stats", "graph stats", "edge count", "code stats"], require: always },
KeywordRule { tool: "add_root",
keywords: &["add_root", "add root", "add scan root"], require: has_path },
KeywordRule { tool: "list_roots",
keywords: &["list_roots", "list roots", "scan roots", "show roots"], require: always },
KeywordRule { tool: "search_code",
keywords: &["search_code", "search code", "find code", "grep", "fts"], require: always },
];
// --- task -----------------------------------------------------------------
const TASK_RULES: [KeywordRule; 9] = [
KeywordRule { tool: "search_tasks",
keywords: &["search_tasks", "search task", "find task", "task search"], require: always },
KeywordRule { tool: "get_task",
keywords: &["get_task", "get task", "task detail"], require: has_id },
KeywordRule { tool: "task_graph",
keywords: &["task_graph", "task graph", "task deps"], require: always },
KeywordRule { tool: "task_stats",
keywords: &["task_stats", "task stats", "task statistics"], require: always },
KeywordRule { tool: "dependency_chain",
keywords: &["dependency_chain", "dep chain", "critical path"], require: always },
KeywordRule { tool: "create_task",
keywords: &["create_task", "create task", "new task", "add task"], require: always },
KeywordRule { tool: "update_task",
keywords: &["update_task", "update task"], require: has_id },
KeywordRule { tool: "add_dependency",
keywords: &["add_dependency", "add dep", "task depends"], require: always },
KeywordRule { tool: "create_milestone",
keywords: &["create_milestone", "create milestone", "new milestone"], require: always },
];
// --- chat -----------------------------------------------------------------
const CHAT_RULES: [KeywordRule; 9] = [
KeywordRule { tool: "search_chat",
keywords: &["search_chat", "search chat", "find in chat", "chat search"], require: always },
KeywordRule { tool: "get_session",
keywords: &["get_session", "chat session", "get session"], require: has_any_id_or_query },
KeywordRule { tool: "list_sessions",
keywords: &["list_sessions", "list sessions", "list chats", "chat history", "my chats"], require: always },
KeywordRule { tool: "chat_stats",
keywords: &["chat_stats", "chat stats", "chat analytics"], require: always },
KeywordRule { tool: "chat_model_usage",
keywords: &["chat_model_usage", "model usage", "token usage"], require: always },
KeywordRule { tool: "start_chat",
keywords: &["start_chat", "new chat", "start chat", "create session"], require: always },
KeywordRule { tool: "save_message",
keywords: &["save_message", "save message", "log message"], require: always },
KeywordRule { tool: "archive_chat",
keywords: &["archive_chat", "archive chat", "close chat"], require: has_any_id_or_query },
KeywordRule { tool: "link_chat",
keywords: &["link_chat", "link chat", "connect chat"], require: always },
];
// --- content, social, cross, curator, search -------------------------------
const CONTENT_RULES: [KeywordRule; 8] = [
KeywordRule { tool: "search_content",
keywords: &["search_content", "search content", "find content", "content search"], require: always },
KeywordRule { tool: "get_asset",
keywords: &["get_asset", "get asset", "asset detail"], require: has_id },
KeywordRule { tool: "content_lineage",
keywords: &["content_lineage", "content lineage", "asset lineage"], require: always },
KeywordRule { tool: "content_stats",
keywords: &["content_stats", "content stats", "content statistics"], require: always },
KeywordRule { tool: "prompt_history",
keywords: &["prompt_history", "prompt history", "prompt log"], require: always },
KeywordRule { tool: "register_asset",
keywords: &["register_asset", "register asset", "new asset", "add asset"], require: always },
KeywordRule { tool: "register_prompt",
keywords: &["register_prompt", "register prompt", "new prompt", "add prompt"], require: always },
KeywordRule { tool: "create_campaign",
keywords: &["create_campaign", "create campaign", "new campaign", "add campaign"], require: always },
];
const SOCIAL_RULES: [KeywordRule; 8] = [
KeywordRule { tool: "search_people",
keywords: &["search_people", "search people", "find people", "people search"], require: always },
KeywordRule { tool: "get_person",
keywords: &["get_person", "get person", "person detail"], require: has_id },
KeywordRule { tool: "relationship_graph",
keywords: &["relationship_graph", "relationship graph", "social graph"], require: always },
KeywordRule { tool: "social_stats",
keywords: &["social_stats", "social stats", "social statistics"], require: always },
KeywordRule { tool: "add_person",
keywords: &["add_person", "add person", "new person"], require: always },
KeywordRule { tool: "add_org",
keywords: &["add_org", "add org", "new org", "add organization"], require: always },
KeywordRule { tool: "log_interaction",
keywords: &["log_interaction", "log interaction", "record interaction"], require: always },
KeywordRule { tool: "link_people",
keywords: &["link_people", "link people", "connect people"], require: always },
];
const CROSS_RULES: [KeywordRule; 8] = [
KeywordRule { tool: "cross_search",
keywords: &["cross_search", "cross search", "search cross", "cross-domain search"], require: always },
KeywordRule { tool: "cross_graph",
keywords: &["cross_graph", "cross graph", "cross-domain graph", "connected across"], require: always },
KeywordRule { tool: "cross_edges",
keywords: &["cross_edges", "cross edges", "inter-domain edges"], require: always },
KeywordRule { tool: "cross_stats",
keywords: &["cross_stats", "cross stats", "cross-domain stats"], require: always },
KeywordRule { tool: "domain_cooccurrence",
keywords: &["domain_cooccurrence", "cooccurrence", "domain cooccurrence"], require: always },
KeywordRule { tool: "cross_link",
keywords: &["cross_link", "link domain", "cross link"], require: always },
KeywordRule { tool: "cross_unlink",
keywords: &["cross_unlink", "unlink domain", "cross unlink"], require: always },
KeywordRule { tool: "cross_auto_link",
keywords: &["cross_auto_link", "auto link", "discover links"], require: always },
];
const CURATOR_RULES: [KeywordRule; 3] = [
KeywordRule { tool: "curator_status",
keywords: &["curator_status", "curator status", "curation status", "curation"], require: always },
KeywordRule { tool: "curator_check",
keywords: &["curator_check", "curator check", "curator dry-run", "curator preview"], require: always },
KeywordRule { tool: "curator_run",
keywords: &["curator_run", "curator run", "run curator"], require: always },
];
const SEARCH_RULES: [KeywordRule; 8] = [
KeywordRule { tool: "search_research",
keywords: &["search_research", "search research", "find research", "past research"], require: always },
KeywordRule { tool: "get_research",
keywords: &["get_research", "get research", "research detail", "show research"], require: has_id },
KeywordRule { tool: "research_sources",
keywords: &["research_sources", "research sources", "sources for research"], require: has_id },
KeywordRule { tool: "research_claims",
keywords: &["research_claims", "research claims", "validated claims", "claims for"], require: has_id },
KeywordRule { tool: "search_stats",
keywords: &["search_stats", "search stats", "research statistics", "research stats"], require: always },
KeywordRule { tool: "run_research",
keywords: &["run_research", "deep research", "research:", "investigate:", "research this"], require: always },
KeywordRule { tool: "stop_research",
keywords: &["stop_research", "stop research", "cancel research"], require: has_id },
KeywordRule { tool: "research_export",
keywords: &["research_export", "export research", "research markdown", "download research"], require: has_id },
];

View file

@ -0,0 +1,20 @@
//! kei-router — NL query to canonical tool-call dispatcher.
//!
//! Constructor Pattern: one cube = one file. Public API:
//! - [`Router::new`] — build with default rules
//! - [`Router::route`] — parse query, return [`RouteResult`]
//! - [`Router::add_dynamic`] — append runtime keyword rules
//!
//! Ported behavior (no ML fallback — depends on CfC, dropped per task spec):
//! * regex-based param extraction (path / limit / depth / id / URI / KV)
//! * keyword-table dispatch, `require` predicate, first-match wins
//! * fallback to `search_code` (if path seen) else `search_knowledge`
pub mod extract;
pub mod keywords;
pub mod router;
pub mod rules;
pub use extract::{extract_params, Extracted};
pub use router::{Method, RouteResult, Router};
pub use rules::{DynRule, KeywordRule};

View file

@ -0,0 +1,35 @@
//! kei-router CLI — print routed tool-call as JSON.
use clap::Parser;
use kei_router::Router;
use std::process::ExitCode;
#[derive(Parser)]
#[command(name = "kei-router", version, about = "Route NL query → tool-call JSON")]
struct Cli {
/// The natural-language query.
query: String,
/// Hint remote-MCP forwarding on fallback (adds _forward=true).
#[arg(long)]
forward: bool,
}
fn main() -> ExitCode {
let cli = Cli::parse();
let router = Router::new();
let result = if cli.forward {
router.route_with_hint(&cli.query)
} else {
router.route(&cli.query)
};
match serde_json::to_string_pretty(&result) {
Ok(s) => {
println!("{}", s);
ExitCode::SUCCESS
}
Err(e) => {
eprintln!("kei-router: json encode failed: {e}");
ExitCode::from(1)
}
}
}

View file

@ -0,0 +1,157 @@
//! Router — holds keyword rules, dispatches queries to tool calls.
use crate::extract::{extract_params, Extracted};
use crate::keywords::default_rules;
use crate::rules::{always, DynRule, KeywordRule};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum Method {
Keyword,
Fallback,
Remote,
}
/// Canonical route outcome.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RouteResult {
pub tool: String,
pub params: BTreeMap<String, serde_json::Value>,
pub confidence: f64,
pub method: Method,
}
/// Router holds the static + dynamic keyword rules.
pub struct Router {
rules: Vec<KeywordRule>,
dynamic: Vec<DynRule>,
}
impl Default for Router {
fn default() -> Self {
Self::new()
}
}
impl Router {
pub fn new() -> Self {
Self {
rules: default_rules(),
dynamic: Vec::new(),
}
}
/// Append user-supplied rules at runtime (domain extension).
pub fn add_dynamic(&mut self, dyn_rules: Vec<DynRule>) {
self.dynamic.extend(dyn_rules);
}
/// Route a natural language query. Always returns a result — falls back to search tools.
pub fn route(&self, query: &str) -> RouteResult {
let ext = extract_params(query);
if let Some(r) = self.keyword_match(&ext) {
return r;
}
if let Some(r) = self.dynamic_match(&ext) {
return r;
}
self.fallback(query, &ext)
}
/// Convenience wrapper — useful for remote MCP forwarders that want a hint.
pub fn route_with_hint(&self, query: &str) -> RouteResult {
let mut r = self.route(query);
if r.method == Method::Fallback {
// Remote-MCP stub: caller may inspect params["_forward"] to decide.
r.params.insert("_forward".into(), serde_json::Value::Bool(true));
}
r
}
fn keyword_match(&self, ext: &Extracted) -> Option<RouteResult> {
for rule in &self.rules {
if !(rule.require)(ext) {
continue;
}
for kw in rule.keywords {
if ext.text_clean.contains(kw) || ext.text.contains(kw) {
return Some(make_route(rule.tool, ext, Method::Keyword, 0.9));
}
}
}
None
}
fn dynamic_match(&self, ext: &Extracted) -> Option<RouteResult> {
for rule in &self.dynamic {
for kw in &rule.keywords {
if ext.text.contains(kw.as_str()) {
return Some(make_route(&rule.tool, ext, Method::Keyword, 0.75));
}
}
}
None
}
fn fallback(&self, query: &str, ext: &Extracted) -> RouteResult {
if !ext.path.is_empty() {
make_route("search_code", ext, Method::Fallback, 0.3)
} else {
let mut params = BTreeMap::new();
params.insert(
"query".into(),
serde_json::Value::String(query.to_string()),
);
RouteResult {
tool: "search_knowledge".into(),
params,
confidence: 0.2,
method: Method::Fallback,
}
}
}
}
fn make_route(tool: &str, ext: &Extracted, method: Method, confidence: f64) -> RouteResult {
RouteResult {
tool: tool.to_string(),
params: merge_params(ext),
confidence,
method,
}
}
fn merge_params(ext: &Extracted) -> BTreeMap<String, serde_json::Value> {
let mut m = BTreeMap::new();
// KV pairs first — typed extraction below takes precedence on collisions
// (e.g. "id=42" → kv["id"]="42" string, but ext.id=42 wins as i64).
for (k, v) in &ext.kv {
m.insert(k.clone(), v.clone().into());
}
if !ext.path.is_empty() {
m.insert("path".into(), ext.path.clone().into());
}
if ext.limit > 0 {
m.insert("limit".into(), ext.limit.into());
}
if ext.depth > 0 {
m.insert("depth".into(), ext.depth.into());
}
if ext.id > 0 {
m.insert("id".into(), ext.id.into());
}
if !ext.query.is_empty() {
m.insert("query".into(), ext.query.clone().into());
}
if !ext.uri.is_empty() {
m.insert("uri".into(), ext.uri.clone().into());
}
m
}
// Silence unused import in some build modes.
#[allow(dead_code)]
fn _always_keep(_e: &Extracted) -> bool {
always(_e)
}

View file

@ -0,0 +1,35 @@
//! Keyword rule type + `require` predicate model.
use crate::extract::Extracted;
/// A dispatch rule: any matching keyword routes to `tool` if `require(extracted)` is true.
#[derive(Clone)]
pub struct KeywordRule {
pub tool: &'static str,
pub keywords: &'static [&'static str],
pub require: fn(&Extracted) -> bool,
}
/// A dynamic (runtime-added) rule — owned strings so caller can build at startup.
#[derive(Clone, Debug)]
pub struct DynRule {
pub tool: String,
pub keywords: Vec<String>,
}
// Predicates mirroring the Go require funcs.
pub fn always(_e: &Extracted) -> bool {
true
}
pub fn has_path(e: &Extracted) -> bool {
!e.path.is_empty()
}
pub fn has_id(e: &Extracted) -> bool {
e.id > 0
}
pub fn has_paths(e: &Extracted) -> bool {
!e.paths.is_empty()
}
pub fn has_any_id_or_query(e: &Extracted) -> bool {
e.id > 0 || !e.query.is_empty()
}

View file

@ -0,0 +1,76 @@
//! kei-router integration tests — mirror LBM router_test.go semantics.
use kei_router::{DynRule, Method, Router};
#[test]
fn exact_match_search_knowledge() {
let r = Router::new();
let out = r.route("search knowledge base for rust async");
assert_eq!(out.tool, "search_knowledge");
assert_eq!(out.method, Method::Keyword);
assert!(out.confidence > 0.7);
}
#[test]
fn fuzzy_match_find_importers_with_path() {
let r = Router::new();
let out = r.route("who imports /src/router.rs");
assert_eq!(out.tool, "find_importers");
assert_eq!(
out.params.get("path").and_then(|v| v.as_str()),
Some("/src/router.rs")
);
}
#[test]
fn no_match_fallback_knowledge() {
let r = Router::new();
let out = r.route("hello this is not a routed query");
assert_eq!(out.tool, "search_knowledge");
assert_eq!(out.method, Method::Fallback);
assert!(out.confidence < 0.3);
}
#[test]
fn no_match_fallback_code_with_path() {
let r = Router::new();
let out = r.route("what happened in /tmp/mystery.rs");
assert_eq!(out.tool, "search_code");
assert_eq!(out.method, Method::Fallback);
}
#[test]
fn confidence_ranking_keyword_above_fallback() {
let r = Router::new();
let kw = r.route("knowledge stats please");
let fb = r.route("asdf zxcv qwer");
assert!(kw.confidence > fb.confidence);
}
#[test]
fn dynamic_rule_addition() {
let mut r = Router::new();
r.add_dynamic(vec![DynRule {
tool: "custom_tool".into(),
keywords: vec!["magic-keyword".into()],
}]);
let out = r.route("please run magic-keyword now");
assert_eq!(out.tool, "custom_tool");
assert_eq!(out.method, Method::Keyword);
}
#[test]
fn remote_mcp_forward_hint() {
let r = Router::new();
let out = r.route_with_hint("completely novel utterance xyz");
assert_eq!(out.method, Method::Fallback);
assert_eq!(out.params.get("_forward"), Some(&serde_json::json!(true)));
}
#[test]
fn id_extraction_for_get_task() {
let r = Router::new();
let out = r.route("get task id=42");
assert_eq!(out.tool, "get_task");
assert_eq!(out.params.get("id").and_then(|v| v.as_i64()), Some(42));
}

View file

@ -0,0 +1,25 @@
[package]
name = "kei-sage"
version = "0.1.0"
edition = "2021"
rust-version = "1.75"
description = "Obsidian-style knowledge graph (SQLite + FTS5). Port of LBM internal/sage."
[[bin]]
name = "kei-sage"
path = "src/main.rs"
[lib]
name = "kei_sage"
path = "src/lib.rs"
[dependencies]
rusqlite = { version = "0.31", features = ["bundled"] }
clap = { version = "4", features = ["derive"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
anyhow = "1"
chrono = { version = "0.4", default-features = false, features = ["clock"] }
[dev-dependencies]
tempfile = "3"

View file

@ -0,0 +1,50 @@
//! BFS traversal over the edges table, depth-limited, deduplicated.
use crate::edges::list_outgoing;
use crate::store::Store;
use crate::types::Related;
use anyhow::Result;
use std::collections::{HashSet, VecDeque};
const MAX_RESULTS: usize = 500;
const MAX_DEPTH: i64 = 5;
pub fn bfs(store: &Store, start: &str, max_depth: i64) -> Result<Vec<Related>> {
let depth = clamp_depth(max_depth);
let mut visited: HashSet<String> = HashSet::new();
visited.insert(start.to_string());
let mut queue: VecDeque<(String, i64)> = VecDeque::new();
queue.push_back((start.to_string(), 0));
let mut out: Vec<Related> = Vec::new();
while let Some((path, d)) = queue.pop_front() {
if out.len() >= MAX_RESULTS {
break;
}
if d >= depth {
continue;
}
for e in list_outgoing(store, &path)? {
if visited.contains(&e.dst_path) || out.len() >= MAX_RESULTS {
continue;
}
visited.insert(e.dst_path.clone());
out.push(Related {
path: e.dst_path.clone(),
edge_type: e.edge_type,
depth: d + 1,
});
queue.push_back((e.dst_path, d + 1));
}
}
Ok(out)
}
fn clamp_depth(d: i64) -> i64 {
if d <= 0 {
2
} else if d > MAX_DEPTH {
MAX_DEPTH
} else {
d
}
}

View file

@ -0,0 +1,62 @@
//! Typed-edge CRUD between vault_paths.
use crate::store::Store;
use crate::types::Edge;
use anyhow::Result;
use chrono::Utc;
use rusqlite::params;
pub fn add_edge(store: &Store, src: &str, dst: &str, edge_type: &str, weight: f64) -> Result<i64> {
let now = Utc::now().timestamp();
store.conn().execute(
"INSERT OR IGNORE INTO edges (src_path, dst_path, edge_type, weight, created_at)
VALUES (?1,?2,?3,?4,?5)",
params![src, dst, edge_type, weight, now],
)?;
Ok(store.conn().last_insert_rowid())
}
pub fn remove_edge(store: &Store, src: &str, dst: &str, edge_type: &str) -> Result<usize> {
let n = store.conn().execute(
"DELETE FROM edges WHERE src_path=?1 AND dst_path=?2 AND edge_type=?3",
params![src, dst, edge_type],
)?;
Ok(n)
}
pub fn list_outgoing(store: &Store, src: &str) -> Result<Vec<Edge>> {
let mut stmt = store.conn().prepare(
"SELECT id, src_path, dst_path, edge_type, weight, created_at
FROM edges WHERE src_path=?1",
)?;
let rows = stmt.query_map(params![src], row_to_edge)?;
let mut out = Vec::new();
for r in rows {
out.push(r?);
}
Ok(out)
}
pub fn list_incoming(store: &Store, dst: &str) -> Result<Vec<Edge>> {
let mut stmt = store.conn().prepare(
"SELECT id, src_path, dst_path, edge_type, weight, created_at
FROM edges WHERE dst_path=?1",
)?;
let rows = stmt.query_map(params![dst], row_to_edge)?;
let mut out = Vec::new();
for r in rows {
out.push(r?);
}
Ok(out)
}
fn row_to_edge(r: &rusqlite::Row) -> rusqlite::Result<Edge> {
Ok(Edge {
id: r.get(0)?,
src_path: r.get(1)?,
dst_path: r.get(2)?,
edge_type: r.get(3)?,
weight: r.get(4)?,
created_at: r.get(5)?,
})
}

View file

@ -0,0 +1,76 @@
//! Obsidian-style vault import: walk a directory, ingest .md files.
//!
//! Minimal subset of LBM internal/sage/import_obsidian.go — we do NOT parse
//! frontmatter here (the upstream parser used multiple helper files). Port
//! of frontmatter/wikilinks parsing is a later milestone; this cube honours
//! the public interface.
use crate::store::Store;
use crate::types::Unit;
use anyhow::{Context, Result};
use std::fs;
use std::path::{Path, PathBuf};
pub struct ImportStats {
pub imported: usize,
pub skipped: usize,
}
pub fn import_vault(store: &Store, root: &Path) -> Result<ImportStats> {
let mut stats = ImportStats { imported: 0, skipped: 0 };
let files = walk_md(root)?;
for path in files {
match ingest_one(store, root, &path) {
Ok(_) => stats.imported += 1,
Err(_) => stats.skipped += 1,
}
}
Ok(stats)
}
fn walk_md(root: &Path) -> Result<Vec<PathBuf>> {
let mut out = Vec::new();
walk_recursive(root, &mut out)?;
Ok(out)
}
fn walk_recursive(dir: &Path, out: &mut Vec<PathBuf>) -> Result<()> {
if !dir.is_dir() {
return Ok(());
}
for entry in fs::read_dir(dir).with_context(|| format!("read_dir {}", dir.display()))? {
let entry = entry?;
let path = entry.path();
if path.is_dir() {
walk_recursive(&path, out)?;
} else if path.extension().and_then(|s| s.to_str()) == Some("md") {
out.push(path);
}
}
Ok(())
}
fn ingest_one(store: &Store, root: &Path, path: &Path) -> Result<()> {
let content = fs::read_to_string(path)?;
let title = path.file_stem()
.and_then(|s| s.to_str())
.unwrap_or("untitled")
.to_string();
let vault_path = path.strip_prefix(root)
.ok()
.and_then(|p| p.to_str())
.unwrap_or(&title)
.to_string();
let unit = Unit {
unit_type: "note".into(),
title,
content,
evidence_grade: "E4".into(),
source_path: path.to_string_lossy().into(),
vault_path,
category: String::new(),
..Default::default()
};
store.add_unit(&unit)?;
Ok(())
}

View file

@ -0,0 +1,15 @@
//! kei-sage — SQLite knowledge-vault with FTS5 + typed edges + BFS + PageRank.
//!
//! Port of LBM internal/sage. Constructor Pattern: one concept per file.
pub mod bfs;
pub mod edges;
pub mod import;
pub mod pagerank;
pub mod schema;
pub mod search;
pub mod store;
pub mod types;
pub use store::Store;
pub use types::{Edge, Related, Unit};

View file

@ -0,0 +1,101 @@
//! kei-sage CLI — import / search / related / rank / add / edit.
use clap::{Parser, Subcommand};
use kei_sage::bfs::bfs;
use kei_sage::edges::add_edge;
use kei_sage::import::import_vault;
use kei_sage::pagerank::pagerank;
use kei_sage::search::fts_search;
use kei_sage::{Store, Unit};
use std::path::PathBuf;
use std::process::ExitCode;
#[derive(Parser)]
#[command(name = "kei-sage", version, about = "Obsidian-style knowledge vault")]
struct Cli {
/// Database path (default: $KEI_VAULT_DB or ~/.claude/sage/vault.sqlite)
#[arg(long)]
db: Option<PathBuf>,
#[command(subcommand)]
cmd: Cmd,
}
#[derive(Subcommand)]
enum Cmd {
Import { vault: PathBuf },
Search { query: String, #[arg(long, default_value_t = 20)] limit: i64 },
Related { key: String, #[arg(long, default_value_t = 2)] depth: i64 },
Rank { #[arg(long, default_value_t = 20)] limit: usize },
Add {
#[arg(long)] title: String,
#[arg(long, default_value = "")] content: String,
#[arg(long, default_value = "")] vault_path: String,
#[arg(long, default_value = "E4")] grade: String,
},
Edit {
id: i64,
#[arg(long)] title: Option<String>,
#[arg(long)] content: Option<String>,
#[arg(long)] grade: Option<String>,
},
Link { src: String, dst: String, #[arg(long, default_value = "related")] edge_type: String },
}
fn db_path(cli_db: Option<PathBuf>) -> PathBuf {
if let Some(p) = cli_db { return p; }
if let Ok(e) = std::env::var("KEI_VAULT_DB") { return PathBuf::from(e); }
let home = std::env::var("HOME").unwrap_or_else(|_| ".".into());
PathBuf::from(home).join(".claude/sage/vault.sqlite")
}
fn run() -> anyhow::Result<()> {
let cli = Cli::parse();
let path = db_path(cli.db);
let store = Store::open(&path)?;
match cli.cmd {
Cmd::Import { vault } => {
let s = import_vault(&store, &vault)?;
println!("imported={} skipped={}", s.imported, s.skipped);
}
Cmd::Search { query, limit } => {
for u in fts_search(&store, &query, limit)? {
println!("{}\t{}\t{}", u.id, u.evidence_grade, u.title);
}
}
Cmd::Related { key, depth } => {
for r in bfs(&store, &key, depth)? {
println!("{}\t{}\t(depth {})", r.edge_type, r.path, r.depth);
}
}
Cmd::Rank { limit } => {
for (p, s) in pagerank(&store)?.into_iter().take(limit) {
println!("{:.6}\t{}", s, p);
}
}
Cmd::Add { title, content, vault_path, grade } => {
let id = store.add_unit(&Unit { title, content, vault_path,
evidence_grade: grade, unit_type: "note".into(), ..Default::default() })?;
println!("{}", id);
}
Cmd::Edit { id, title, content, grade } => {
let mut u = store.get_unit(id)?.ok_or_else(|| anyhow::anyhow!("id {id} not found"))?;
if let Some(t) = title { u.title = t; }
if let Some(c) = content { u.content = c; }
if let Some(g) = grade { u.evidence_grade = g; }
store.update_unit(&u)?;
println!("updated {}", id);
}
Cmd::Link { src, dst, edge_type } => {
add_edge(&store, &src, &dst, &edge_type, 1.0)?;
println!("linked {} -> {}", src, dst);
}
}
Ok(())
}
fn main() -> ExitCode {
match run() {
Ok(()) => ExitCode::SUCCESS,
Err(e) => { eprintln!("kei-sage: {e:#}"); ExitCode::from(1) }
}
}

View file

@ -0,0 +1,60 @@
//! PageRank — power-iteration, 50 iterations, d=0.85. Operates on the edges table.
use crate::store::Store;
use anyhow::Result;
use std::collections::HashMap;
const DAMPING: f64 = 0.85;
const ITERATIONS: usize = 50;
/// Compute PageRank over the edges table. Returns [(path, score)] sorted desc.
pub fn pagerank(store: &Store) -> Result<Vec<(String, f64)>> {
let (nodes, out_edges) = collect_graph(store)?;
if nodes.is_empty() {
return Ok(Vec::new());
}
let mut rank: HashMap<String, f64> = nodes.iter()
.map(|n| (n.clone(), 1.0 / nodes.len() as f64)).collect();
for _ in 0..ITERATIONS {
rank = one_iteration(&nodes, &out_edges, &rank);
}
let mut out: Vec<(String, f64)> = rank.into_iter().collect();
out.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
Ok(out)
}
fn collect_graph(store: &Store) -> Result<(Vec<String>, HashMap<String, Vec<String>>)> {
let mut stmt = store.conn().prepare("SELECT src_path, dst_path FROM edges")?;
let rows = stmt.query_map([], |r| Ok((r.get::<_, String>(0)?, r.get::<_, String>(1)?)))?;
let mut nodes: std::collections::HashSet<String> = std::collections::HashSet::new();
let mut out_edges: HashMap<String, Vec<String>> = HashMap::new();
for row in rows {
let (src, dst) = row?;
nodes.insert(src.clone());
nodes.insert(dst.clone());
out_edges.entry(src).or_default().push(dst);
}
Ok((nodes.into_iter().collect(), out_edges))
}
fn one_iteration(
nodes: &[String],
out_edges: &HashMap<String, Vec<String>>,
prev: &HashMap<String, f64>,
) -> HashMap<String, f64> {
let n = nodes.len() as f64;
let base = (1.0 - DAMPING) / n;
let mut next: HashMap<String, f64> = nodes.iter().map(|k| (k.clone(), base)).collect();
for (src, dsts) in out_edges {
if dsts.is_empty() {
continue;
}
let share = DAMPING * prev.get(src).copied().unwrap_or(0.0) / dsts.len() as f64;
for dst in dsts {
if let Some(slot) = next.get_mut(dst) {
*slot += share;
}
}
}
next
}

View file

@ -0,0 +1,57 @@
//! SQLite schema for knowledge-vault. Port of LBM internal/sage/vault_schema.go.
use rusqlite::{Connection, Result};
/// Apply schema + FTS5 virtual table. Idempotent.
pub fn create_schema(conn: &Connection) -> Result<()> {
conn.execute_batch(
r#"
CREATE TABLE IF NOT EXISTS knowledge_units (
id INTEGER PRIMARY KEY,
unit_type TEXT NOT NULL,
title TEXT NOT NULL,
content TEXT DEFAULT '',
evidence_grade TEXT DEFAULT '',
source_path TEXT DEFAULT '',
vault_path TEXT DEFAULT '',
category TEXT DEFAULT '',
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_ku_type ON knowledge_units(unit_type);
CREATE UNIQUE INDEX IF NOT EXISTS idx_ku_vault
ON knowledge_units(vault_path) WHERE vault_path != '';
CREATE INDEX IF NOT EXISTS idx_ku_grade ON knowledge_units(evidence_grade);
CREATE TABLE IF NOT EXISTS tags (
id INTEGER PRIMARY KEY,
name TEXT UNIQUE NOT NULL
);
CREATE TABLE IF NOT EXISTS unit_tags (
unit_id INTEGER NOT NULL REFERENCES knowledge_units(id) ON DELETE CASCADE,
tag_id INTEGER NOT NULL REFERENCES tags(id) ON DELETE CASCADE,
PRIMARY KEY (unit_id, tag_id)
);
CREATE TABLE IF NOT EXISTS edges (
id INTEGER PRIMARY KEY,
src_path TEXT NOT NULL,
dst_path TEXT NOT NULL,
edge_type TEXT NOT NULL,
weight REAL DEFAULT 1.0,
created_at INTEGER NOT NULL,
UNIQUE(src_path, dst_path, edge_type)
);
CREATE INDEX IF NOT EXISTS idx_sage_edges_src ON edges(src_path);
CREATE INDEX IF NOT EXISTS idx_sage_edges_dst ON edges(dst_path);
"#,
)?;
conn.execute_batch(
r#"
CREATE VIRTUAL TABLE IF NOT EXISTS fts_knowledge
USING fts5(unit_id UNINDEXED, title, content, tokenize='porter unicode61');
"#,
)?;
Ok(())
}

View file

@ -0,0 +1,39 @@
//! FTS5 search over knowledge_units.
use crate::store::Store;
use crate::types::Unit;
use anyhow::Result;
use rusqlite::params;
/// Full-text search. Returns matching Units ordered by SQLite FTS5 rank.
pub fn fts_search(store: &Store, query: &str, limit: i64) -> Result<Vec<Unit>> {
let lim = if limit <= 0 { 20 } else { limit };
let mut stmt = store.conn().prepare(
"SELECT k.id, k.unit_type, k.title, k.content, k.evidence_grade,
k.source_path, k.vault_path, k.category, k.created_at, k.updated_at
FROM fts_knowledge f
JOIN knowledge_units k ON k.id = f.unit_id
WHERE fts_knowledge MATCH ?1
ORDER BY rank
LIMIT ?2",
)?;
let rows = stmt.query_map(params![query, lim], |r| {
Ok(Unit {
id: r.get(0)?,
unit_type: r.get(1)?,
title: r.get(2)?,
content: r.get(3)?,
evidence_grade: r.get(4)?,
source_path: r.get(5)?,
vault_path: r.get(6)?,
category: r.get(7)?,
created_at: r.get(8)?,
updated_at: r.get(9)?,
})
})?;
let mut out = Vec::new();
for row in rows {
out.push(row?);
}
Ok(out)
}

View file

@ -0,0 +1,111 @@
//! Knowledge-unit CRUD + FTS indexer.
use crate::schema::create_schema;
use crate::types::Unit;
use anyhow::{Context, Result};
use chrono::Utc;
use rusqlite::{params, Connection};
use std::path::Path;
pub struct Store {
conn: Connection,
}
impl Store {
pub fn open(path: &Path) -> Result<Self> {
if let Some(parent) = path.parent() {
let _ = std::fs::create_dir_all(parent);
}
let conn = Connection::open(path).context("open sqlite")?;
conn.pragma_update(None, "journal_mode", "WAL").ok();
create_schema(&conn)?;
Ok(Self { conn })
}
pub fn open_memory() -> Result<Self> {
let conn = Connection::open_in_memory()?;
create_schema(&conn)?;
Ok(Self { conn })
}
pub fn conn(&self) -> &Connection {
&self.conn
}
/// Insert a new knowledge unit. Indexes title+content into FTS5. Idempotent by vault_path.
pub fn add_unit(&self, unit: &Unit) -> Result<i64> {
let now = Utc::now().timestamp();
let created = if unit.created_at == 0 { now } else { unit.created_at };
self.conn.execute(
"INSERT OR REPLACE INTO knowledge_units
(unit_type, title, content, evidence_grade, source_path,
vault_path, category, created_at, updated_at)
VALUES (?1,?2,?3,?4,?5,?6,?7,?8,?9)",
params![unit.unit_type, unit.title, unit.content, unit.evidence_grade,
unit.source_path, unit.vault_path, unit.category, created, now],
)?;
let id = self.conn.last_insert_rowid();
self.reindex_fts(id, &unit.title, &unit.content)?;
Ok(id)
}
pub fn get_unit(&self, id: i64) -> Result<Option<Unit>> {
let mut stmt = self.conn.prepare(
"SELECT id, unit_type, title, content, evidence_grade, source_path,
vault_path, category, created_at, updated_at
FROM knowledge_units WHERE id=?1",
)?;
let mut rows = stmt.query(params![id])?;
if let Some(r) = rows.next()? {
return Ok(Some(row_to_unit(r)?));
}
Ok(None)
}
pub fn update_unit(&self, unit: &Unit) -> Result<()> {
let now = Utc::now().timestamp();
self.conn.execute(
"UPDATE knowledge_units SET title=?1, content=?2, evidence_grade=?3,
category=?4, updated_at=?5 WHERE id=?6",
params![unit.title, unit.content, unit.evidence_grade,
unit.category, now, unit.id],
)?;
self.reindex_fts(unit.id, &unit.title, &unit.content)?;
Ok(())
}
pub fn delete_unit(&self, id: i64) -> Result<()> {
self.conn.execute("DELETE FROM fts_knowledge WHERE unit_id=?1", params![id])?;
self.conn.execute("DELETE FROM knowledge_units WHERE id=?1", params![id])?;
Ok(())
}
pub fn count_units(&self) -> Result<i64> {
Ok(self.conn.query_row(
"SELECT COUNT(*) FROM knowledge_units", [], |r| r.get(0))?)
}
fn reindex_fts(&self, id: i64, title: &str, content: &str) -> Result<()> {
self.conn.execute("DELETE FROM fts_knowledge WHERE unit_id=?1", params![id])?;
self.conn.execute(
"INSERT INTO fts_knowledge (unit_id, title, content) VALUES (?1,?2,?3)",
params![id, title, content],
)?;
Ok(())
}
}
fn row_to_unit(r: &rusqlite::Row) -> rusqlite::Result<Unit> {
Ok(Unit {
id: r.get(0)?,
unit_type: r.get(1)?,
title: r.get(2)?,
content: r.get(3)?,
evidence_grade: r.get(4)?,
source_path: r.get(5)?,
vault_path: r.get(6)?,
category: r.get(7)?,
created_at: r.get(8)?,
updated_at: r.get(9)?,
})
}

View file

@ -0,0 +1,34 @@
//! Shared value types for knowledge units + edges + BFS results.
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Unit {
pub id: i64,
pub unit_type: String,
pub title: String,
pub content: String,
pub evidence_grade: String,
pub source_path: String,
pub vault_path: String,
pub category: String,
pub created_at: i64,
pub updated_at: i64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Edge {
pub id: i64,
pub src_path: String,
pub dst_path: String,
pub edge_type: String,
pub weight: f64,
pub created_at: i64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Related {
pub path: String,
pub edge_type: String,
pub depth: i64,
}

View file

@ -0,0 +1,110 @@
//! kei-sage integration tests.
use kei_sage::bfs::bfs;
use kei_sage::edges::{add_edge, list_outgoing};
use kei_sage::import::import_vault;
use kei_sage::pagerank::pagerank;
use kei_sage::search::fts_search;
use kei_sage::{Store, Unit};
use std::fs;
use tempfile::tempdir;
fn mkstore() -> Store { Store::open_memory().unwrap() }
fn mkunit(title: &str, body: &str, vault: &str) -> Unit {
Unit {
unit_type: "note".into(), title: title.into(), content: body.into(),
evidence_grade: "E2".into(), vault_path: vault.into(),
..Default::default()
}
}
#[test]
fn crud_roundtrip() {
let s = mkstore();
let id = s.add_unit(&mkunit("hello", "world", "a.md")).unwrap();
assert!(id > 0);
let u = s.get_unit(id).unwrap().unwrap();
assert_eq!(u.title, "hello");
s.delete_unit(id).unwrap();
assert!(s.get_unit(id).unwrap().is_none());
}
#[test]
fn fts_search_matches() {
let s = mkstore();
s.add_unit(&mkunit("rust async", "tokio runtime details", "a.md")).unwrap();
s.add_unit(&mkunit("python sync", "flask wsgi server", "b.md")).unwrap();
let hits = fts_search(&s, "tokio", 10).unwrap();
assert_eq!(hits.len(), 1);
assert_eq!(hits[0].title, "rust async");
}
#[test]
fn bfs_depth_limit() {
let s = mkstore();
add_edge(&s, "a", "b", "rel", 1.0).unwrap();
add_edge(&s, "b", "c", "rel", 1.0).unwrap();
add_edge(&s, "c", "d", "rel", 1.0).unwrap();
let out = bfs(&s, "a", 2).unwrap();
let paths: Vec<&str> = out.iter().map(|r| r.path.as_str()).collect();
assert!(paths.contains(&"b"));
assert!(paths.contains(&"c"));
assert!(!paths.contains(&"d"));
}
#[test]
fn pagerank_orders_by_popularity() {
let s = mkstore();
add_edge(&s, "a", "hub", "rel", 1.0).unwrap();
add_edge(&s, "b", "hub", "rel", 1.0).unwrap();
add_edge(&s, "c", "hub", "rel", 1.0).unwrap();
add_edge(&s, "d", "hub", "rel", 1.0).unwrap();
add_edge(&s, "e", "hub", "rel", 1.0).unwrap();
let ranks = pagerank(&s).unwrap();
assert_eq!(ranks[0].0, "hub");
}
#[test]
fn edges_crud() {
let s = mkstore();
let id = add_edge(&s, "x", "y", "cites", 0.8).unwrap();
assert!(id > 0);
let out = list_outgoing(&s, "x").unwrap();
assert_eq!(out.len(), 1);
assert_eq!(out[0].dst_path, "y");
}
#[test]
fn import_idempotency() {
let tmp = tempdir().unwrap();
let p = tmp.path().join("one.md");
fs::write(&p, "# title one\nhello").unwrap();
let s = mkstore();
let first = import_vault(&s, tmp.path()).unwrap();
let second = import_vault(&s, tmp.path()).unwrap();
assert_eq!(first.imported, 1);
assert_eq!(second.imported, 1);
assert_eq!(s.count_units().unwrap(), 1);
}
#[test]
fn edges_cross_reference_validates() {
let s = mkstore();
s.add_unit(&mkunit("note a", "", "a.md")).unwrap();
s.add_unit(&mkunit("note b", "", "b.md")).unwrap();
add_edge(&s, "a.md", "b.md", "refs", 1.0).unwrap();
let out = list_outgoing(&s, "a.md").unwrap();
assert_eq!(out.len(), 1);
}
#[test]
fn fts5_respects_limit() {
let s = mkstore();
for i in 0..25 {
let t = format!("rust note {i}");
s.add_unit(&mkunit(&t, "rust rust rust", &format!("n{i}.md"))).unwrap();
}
let hits = fts_search(&s, "rust", 5).unwrap();
assert_eq!(hits.len(), 5);
}

View file

@ -0,0 +1,25 @@
[package]
name = "kei-search-core"
version = "0.1.0"
edition = "2021"
rust-version = "1.75"
description = "3-wave deep research scaffolding with budget cap. Port of LBM internal/search (fetch stubbed)."
[[bin]]
name = "kei-search-core"
path = "src/main.rs"
[lib]
name = "kei_search_core"
path = "src/lib.rs"
[dependencies]
rusqlite = { version = "0.31", features = ["bundled"] }
clap = { version = "4", features = ["derive"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
anyhow = "1"
chrono = { version = "0.4", default-features = false, features = ["clock"] }
[dev-dependencies]
tempfile = "3"

View file

@ -0,0 +1,34 @@
//! Budget tracker — all costs in microcents (1 USD = 1_000_000 mc).
use anyhow::{anyhow, Result};
#[derive(Debug, Clone)]
pub struct Budget {
cap_mc: i64,
spent_mc: i64,
stopped: bool,
}
impl Budget {
pub fn new(cap_mc: i64) -> Self {
Self { cap_mc, spent_mc: 0, stopped: false }
}
/// Record a cost; returns error if this push would exceed the cap.
pub fn charge(&mut self, mc: i64) -> Result<()> {
if self.stopped {
return Err(anyhow!("budget stopped"));
}
if self.spent_mc + mc > self.cap_mc {
return Err(anyhow!(
"budget exceeded: spent={} cap={}", self.spent_mc + mc, self.cap_mc));
}
self.spent_mc += mc;
Ok(())
}
pub fn spent(&self) -> i64 { self.spent_mc }
pub fn remaining(&self) -> i64 { self.cap_mc - self.spent_mc }
pub fn stop(&mut self) { self.stopped = true; }
pub fn is_stopped(&self) -> bool { self.stopped }
}

View file

@ -0,0 +1,40 @@
//! Export research → markdown / JSON.
use crate::store::ResearchStore;
use anyhow::{anyhow, Result};
use serde_json::json;
pub enum Format {
Markdown,
Json,
}
pub fn export(store: &ResearchStore, id: i64, fmt: Format) -> Result<String> {
let r = store.get_research(id)?.ok_or_else(|| anyhow!("research {id} missing"))?;
let claims = store.claims_for(id)?;
match fmt {
Format::Markdown => {
let mut md = String::new();
md.push_str(&format!("# Research {}\n\n", r.id));
md.push_str(&format!("**Query:** {}\n\n", r.query_original));
md.push_str(&format!("**Status:** {}\n", r.status));
md.push_str(&format!("**Cost:** {} mc\n\n", r.total_cost_mc));
md.push_str("## Claims\n\n");
for c in claims {
md.push_str(&format!("- [{}] {} (consensus={:.2})\n",
c.grade, c.claim_text, c.consensus));
}
Ok(md)
}
Format::Json => {
let val = json!({
"id": r.id,
"query": r.query_original,
"status": r.status,
"cost_mc": r.total_cost_mc,
"claims": claims,
});
Ok(serde_json::to_string_pretty(&val)?)
}
}
}

View file

@ -0,0 +1,23 @@
//! Source fetcher trait — frozen interface, default impl is a no-op stub.
//!
//! Actual WebFetch/WebSearch integration is out-of-scope for v0.14 part A.
//! Later milestones plug real providers (anthropic-websearch, SerpAPI, etc.).
use crate::types::Source;
/// Implement this trait to integrate a live search provider.
pub trait SourceFetcher {
/// Fetch sources for `claim`. Returns (source, cost_microcents).
/// Cost is real — the budget is charged by the pipeline, not by impl.
fn fetch(&self, claim: &str) -> (Vec<Source>, i64);
}
/// Default stub — returns empty. Frozen interface, no runtime side-effects.
pub struct StubFetcher;
impl SourceFetcher for StubFetcher {
fn fetch(&self, _claim: &str) -> (Vec<Source>, i64) {
// TODO(v0.15): wire to real websearch. Kept as stub per v0.14 spec.
(Vec::new(), 0)
}
}

View file

@ -0,0 +1,22 @@
//! kei-search-core — 3-wave deep research engine, budget-capped.
//!
//! Waves:
//! 0 — claim extraction from prompt
//! 1 — per-claim source hunt (WebFetch stubbed behind [`SourceFetcher`] trait)
//! 2 — cross-validation + consensus scoring
//!
//! Port of LBM internal/search. The actual fetch is a trait the caller
//! supplies. Default implementation returns empty (frozen interface, todo!()
//! reflects unimplemented runtime).
pub mod budget;
pub mod export;
pub mod fetch;
pub mod pipeline;
pub mod schema;
pub mod store;
pub mod types;
pub use pipeline::run_research;
pub use store::ResearchStore;
pub use types::{Claim, Research, Source};

View file

@ -0,0 +1,61 @@
//! kei-search-core CLI.
use clap::{Parser, Subcommand, ValueEnum};
use kei_search_core::export::{export, Format};
use kei_search_core::fetch::StubFetcher;
use kei_search_core::pipeline::run_research;
use kei_search_core::ResearchStore;
use std::path::PathBuf;
use std::process::ExitCode;
#[derive(Parser)]
#[command(name = "kei-search-core", version)]
struct Cli {
#[arg(long)] db: Option<PathBuf>,
#[command(subcommand)] cmd: Cmd,
}
#[derive(Subcommand)]
enum Cmd {
Run { prompt: String,
#[arg(long, default_value_t = 1_000_000)] budget: i64 }, // 1 USD
Stop { id: i64 },
Export { id: i64, #[arg(long, value_enum, default_value_t = Fmt::Md)] format: Fmt },
}
#[derive(Clone, Copy, ValueEnum)]
enum Fmt { Md, Json }
fn db_path(o: Option<PathBuf>) -> PathBuf {
if let Some(p) = o { return p; }
if let Ok(e) = std::env::var("KEI_SEARCH_DB") { return PathBuf::from(e); }
let home = std::env::var("HOME").unwrap_or_else(|_| ".".into());
PathBuf::from(home).join(".claude/search/research.sqlite")
}
fn run() -> anyhow::Result<()> {
let cli = Cli::parse();
let s = ResearchStore::open(&db_path(cli.db))?;
match cli.cmd {
Cmd::Run { prompt, budget } => {
let id = run_research(&s, &StubFetcher, &prompt, budget)?;
println!("{}", id);
}
Cmd::Stop { id } => {
s.set_status(id, "stopped")?;
println!("stopped {}", id);
}
Cmd::Export { id, format } => {
let f = match format { Fmt::Md => Format::Markdown, Fmt::Json => Format::Json };
println!("{}", export(&s, id, f)?);
}
}
Ok(())
}
fn main() -> ExitCode {
match run() {
Ok(()) => ExitCode::SUCCESS,
Err(e) => { eprintln!("kei-search-core: {e:#}"); ExitCode::from(1) }
}
}

View file

@ -0,0 +1,90 @@
//! 3-wave research runner.
//!
//! Wave 0: split prompt into claims (naive split on `.`; real NLU later).
//! Wave 1: for each claim, fetch sources via [`SourceFetcher`].
//! Wave 2: score consensus per claim from sources (majority = higher grade).
use crate::budget::Budget;
use crate::fetch::SourceFetcher;
use crate::store::ResearchStore;
use crate::types::{Claim, Source};
use anyhow::Result;
const WAVE1_COST_PER_CLAIM_MC: i64 = 100; // 0.01 USD per claim
const WAVE2_COST_MC: i64 = 50;
pub fn run_research(
store: &ResearchStore,
fetcher: &dyn SourceFetcher,
prompt: &str,
budget_mc: i64,
) -> Result<i64> {
let research_id = store.create_research(prompt)?;
let mut budget = Budget::new(budget_mc);
let claims_text = wave_0_extract_claims(prompt);
if let Err(e) = wave_1_fetch(store, fetcher, research_id, &claims_text, &mut budget) {
store.set_status(research_id, "failed")?;
return Err(e);
}
if let Err(e) = wave_2_consensus(store, research_id, &mut budget) {
store.set_status(research_id, "failed")?;
return Err(e);
}
store.set_cost(research_id, budget.spent())?;
store.set_status(research_id, "completed")?;
Ok(research_id)
}
fn wave_0_extract_claims(prompt: &str) -> Vec<String> {
prompt
.split(|c: char| c == '.' || c == '?' || c == '\n')
.map(|s| s.trim().to_string())
.filter(|s| s.len() > 4)
.collect()
}
fn wave_1_fetch(
store: &ResearchStore,
fetcher: &dyn SourceFetcher,
rid: i64,
claims: &[String],
budget: &mut Budget,
) -> Result<()> {
for c in claims {
budget.charge(WAVE1_COST_PER_CLAIM_MC)?;
let (srcs, fetch_cost) = fetcher.fetch(c);
if fetch_cost > 0 {
budget.charge(fetch_cost)?;
}
for s in srcs {
store.add_source(&Source { research_id: rid, ..s })?;
}
store.add_claim(&Claim {
research_id: rid,
claim_text: c.clone(),
..Default::default()
})?;
}
Ok(())
}
fn wave_2_consensus(store: &ResearchStore, rid: i64, budget: &mut Budget) -> Result<()> {
budget.charge(WAVE2_COST_MC)?;
let claims = store.claims_for(rid)?;
for c in claims {
let support = 0.5;
let contradict = 0.0;
let consensus = support - contradict;
let grade = grade_from_consensus(consensus);
store.conn().execute(
"UPDATE claims SET support=?1, contradict=?2, consensus=?3, grade=?4
WHERE id=?5",
rusqlite::params![support, contradict, consensus, grade, c.id],
)?;
}
Ok(())
}
fn grade_from_consensus(c: f64) -> &'static str {
if c >= 0.8 { "E2" } else if c >= 0.5 { "E4" } else { "E6" }
}

View file

@ -0,0 +1,42 @@
use rusqlite::{Connection, Result};
pub fn create_schema(conn: &Connection) -> Result<()> {
conn.execute_batch(r#"
CREATE TABLE IF NOT EXISTS researches (
id INTEGER PRIMARY KEY,
query_original TEXT NOT NULL,
status TEXT NOT NULL DEFAULT 'pending',
result_markdown TEXT DEFAULT '',
total_cost_mc INTEGER DEFAULT 0,
created_at INTEGER NOT NULL,
completed_at INTEGER DEFAULT 0
);
CREATE INDEX IF NOT EXISTS idx_res_status ON researches(status);
CREATE TABLE IF NOT EXISTS sources (
id INTEGER PRIMARY KEY,
research_id INTEGER NOT NULL REFERENCES researches(id),
url TEXT NOT NULL,
title TEXT DEFAULT '',
content TEXT DEFAULT '',
provider TEXT DEFAULT '',
domain TEXT DEFAULT '',
relevance_score REAL DEFAULT 0.0,
created_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_src_research ON sources(research_id);
CREATE TABLE IF NOT EXISTS claims (
id INTEGER PRIMARY KEY,
research_id INTEGER NOT NULL REFERENCES researches(id),
claim_text TEXT NOT NULL,
support REAL DEFAULT 0.0,
contradict REAL DEFAULT 0.0,
consensus REAL DEFAULT 0.0,
grade TEXT DEFAULT 'E6',
created_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_claim_research ON claims(research_id);
"#)?;
Ok(())
}

View file

@ -0,0 +1,122 @@
use crate::schema::create_schema;
use crate::types::{Claim, Research, Source};
use anyhow::{Context, Result};
use chrono::Utc;
use rusqlite::{params, Connection};
use std::path::Path;
pub struct ResearchStore {
conn: Connection,
}
impl ResearchStore {
pub fn open(path: &Path) -> Result<Self> {
if let Some(parent) = path.parent() {
let _ = std::fs::create_dir_all(parent);
}
let conn = Connection::open(path).context("open sqlite")?;
conn.pragma_update(None, "journal_mode", "WAL").ok();
create_schema(&conn)?;
Ok(Self { conn })
}
pub fn open_memory() -> Result<Self> {
let conn = Connection::open_in_memory()?;
create_schema(&conn)?;
Ok(Self { conn })
}
pub fn conn(&self) -> &Connection { &self.conn }
pub fn create_research(&self, query: &str) -> Result<i64> {
let now = Utc::now().timestamp();
self.conn.execute(
"INSERT INTO researches (query_original, status, created_at)
VALUES (?1, 'running', ?2)",
params![query, now],
)?;
Ok(self.conn.last_insert_rowid())
}
pub fn set_status(&self, id: i64, status: &str) -> Result<()> {
let now = Utc::now().timestamp();
self.conn.execute(
"UPDATE researches SET status=?1, completed_at=?2 WHERE id=?3",
params![status, now, id],
)?;
Ok(())
}
pub fn set_cost(&self, id: i64, mc: i64) -> Result<()> {
self.conn.execute(
"UPDATE researches SET total_cost_mc=?1 WHERE id=?2",
params![mc, id],
)?;
Ok(())
}
pub fn set_markdown(&self, id: i64, md: &str) -> Result<()> {
self.conn.execute(
"UPDATE researches SET result_markdown=?1 WHERE id=?2",
params![md, id],
)?;
Ok(())
}
pub fn get_research(&self, id: i64) -> Result<Option<Research>> {
let mut stmt = self.conn.prepare(
"SELECT id, query_original, status, result_markdown, total_cost_mc,
created_at, completed_at FROM researches WHERE id=?1",
)?;
let mut rows = stmt.query(params![id])?;
if let Some(r) = rows.next()? {
return Ok(Some(Research {
id: r.get(0)?, query_original: r.get(1)?, status: r.get(2)?,
result_markdown: r.get(3)?, total_cost_mc: r.get(4)?,
created_at: r.get(5)?, completed_at: r.get(6)?,
}));
}
Ok(None)
}
pub fn add_source(&self, s: &Source) -> Result<i64> {
let now = Utc::now().timestamp();
self.conn.execute(
"INSERT INTO sources (research_id, url, title, content, provider,
domain, relevance_score, created_at)
VALUES (?1,?2,?3,?4,?5,?6,?7,?8)",
params![s.research_id, s.url, s.title, s.content, s.provider,
s.domain, s.relevance_score, now],
)?;
Ok(self.conn.last_insert_rowid())
}
pub fn add_claim(&self, c: &Claim) -> Result<i64> {
let now = Utc::now().timestamp();
self.conn.execute(
"INSERT INTO claims (research_id, claim_text, support, contradict,
consensus, grade, created_at)
VALUES (?1,?2,?3,?4,?5,?6,?7)",
params![c.research_id, c.claim_text, c.support, c.contradict,
c.consensus, c.grade, now],
)?;
Ok(self.conn.last_insert_rowid())
}
pub fn claims_for(&self, research_id: i64) -> Result<Vec<Claim>> {
let mut stmt = self.conn.prepare(
"SELECT id, research_id, claim_text, support, contradict,
consensus, grade, created_at FROM claims WHERE research_id=?1"
)?;
let rows = stmt.query_map(params![research_id], |r| {
Ok(Claim {
id: r.get(0)?, research_id: r.get(1)?, claim_text: r.get(2)?,
support: r.get(3)?, contradict: r.get(4)?, consensus: r.get(5)?,
grade: r.get(6)?, created_at: r.get(7)?,
})
})?;
let mut out = Vec::new();
for r in rows { out.push(r?); }
Ok(out)
}
}

View file

@ -0,0 +1,37 @@
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Research {
pub id: i64,
pub query_original: String,
pub status: String,
pub result_markdown: String,
pub total_cost_mc: i64,
pub created_at: i64,
pub completed_at: i64,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Source {
pub id: i64,
pub research_id: i64,
pub url: String,
pub title: String,
pub content: String,
pub provider: String,
pub domain: String,
pub relevance_score: f64,
pub created_at: i64,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Claim {
pub id: i64,
pub research_id: i64,
pub claim_text: String,
pub support: f64,
pub contradict: f64,
pub consensus: f64,
pub grade: String,
pub created_at: i64,
}

View file

@ -0,0 +1,80 @@
use kei_search_core::budget::Budget;
use kei_search_core::export::{export, Format};
use kei_search_core::fetch::{SourceFetcher, StubFetcher};
use kei_search_core::pipeline::run_research;
use kei_search_core::types::Source;
use kei_search_core::ResearchStore;
fn mk() -> ResearchStore { ResearchStore::open_memory().unwrap() }
struct FakeFetcher;
impl SourceFetcher for FakeFetcher {
fn fetch(&self, claim: &str) -> (Vec<Source>, i64) {
(vec![Source {
url: "https://example.test".into(),
title: format!("source for: {claim}"),
content: "body".into(),
provider: "fake".into(),
domain: "example.test".into(),
relevance_score: 0.8,
..Default::default()
}], 10)
}
}
#[test]
fn budget_enforcement() {
let mut b = Budget::new(100);
b.charge(50).unwrap();
b.charge(40).unwrap();
assert!(b.charge(20).is_err(), "must reject overspend");
}
#[test]
fn wave_progression_creates_research() {
let s = mk();
let id = run_research(&s, &FakeFetcher,
"Rust is memory-safe. Python is dynamic.", 10_000).unwrap();
let r = s.get_research(id).unwrap().unwrap();
assert_eq!(r.status, "completed");
assert!(r.total_cost_mc > 0);
assert!(s.claims_for(id).unwrap().len() >= 2);
}
#[test]
fn consensus_scoring_applies_grade() {
let s = mk();
let id = run_research(&s, &FakeFetcher, "One claim here.", 10_000).unwrap();
let claims = s.claims_for(id).unwrap();
assert!(!claims.is_empty());
assert!(!claims[0].grade.is_empty());
}
#[test]
fn export_markdown_and_json() {
let s = mk();
let id = run_research(&s, &FakeFetcher, "Claim A. Claim B.", 10_000).unwrap();
let md = export(&s, id, Format::Markdown).unwrap();
assert!(md.contains("# Research"));
let js = export(&s, id, Format::Json).unwrap();
let parsed: serde_json::Value = serde_json::from_str(&js).unwrap();
assert!(parsed.get("claims").is_some());
}
#[test]
fn stop_mid_run_marks_status() {
let s = mk();
let id = run_research(&s, &StubFetcher, "x. y.", 10_000).unwrap();
s.set_status(id, "stopped").unwrap();
let r = s.get_research(id).unwrap().unwrap();
assert_eq!(r.status, "stopped");
}
#[test]
fn budget_exhausted_rejects_run() {
let s = mk();
// 3 claims × 100mc + 50mc wave2 = 350mc; budget 100 → must overspend.
let err = run_research(&s, &StubFetcher,
"alpha claim one. beta claim two. gamma claim three.", 100);
assert!(err.is_err(), "small budget vs 3 claims must overspend");
}

View file

@ -0,0 +1,25 @@
[package]
name = "kei-social-store"
version = "0.1.0"
edition = "2021"
rust-version = "1.75"
description = "People + interaction CRM (lite). Port of LBM internal/social."
[[bin]]
name = "kei-social-store"
path = "src/main.rs"
[lib]
name = "kei_social_store"
path = "src/lib.rs"
[dependencies]
rusqlite = { version = "0.31", features = ["bundled"] }
clap = { version = "4", features = ["derive"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
anyhow = "1"
chrono = { version = "0.4", default-features = false, features = ["clock"] }
[dev-dependencies]
tempfile = "3"

View file

@ -0,0 +1,31 @@
//! Relationship graph — who interacted with whom, grouped by channel.
use crate::store::Store;
use anyhow::Result;
use serde::Serialize;
#[derive(Debug, Clone, Serialize)]
pub struct Pair {
pub person_id: i64,
pub target_id: i64,
pub channel: String,
pub count: i64,
}
pub fn relationship_graph(store: &Store) -> Result<Vec<Pair>> {
let mut stmt = store.conn().prepare(
"SELECT person_id, target_id, channel, COUNT(*) FROM interactions
WHERE target_id > 0 GROUP BY person_id, target_id, channel",
)?;
let rows = stmt.query_map([], |r| {
Ok(Pair {
person_id: r.get(0)?,
target_id: r.get(1)?,
channel: r.get(2)?,
count: r.get(3)?,
})
})?;
let mut out = Vec::new();
for r in rows { out.push(r?); }
Ok(out)
}

View file

@ -0,0 +1,47 @@
use crate::store::Store;
use anyhow::Result;
use chrono::Utc;
use rusqlite::params;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Interaction {
pub id: i64,
pub person_id: i64,
pub target_id: i64,
pub interaction_type: String,
pub channel: String,
pub content: String,
pub timestamp: i64,
}
pub fn log_interaction(store: &Store, i: &Interaction) -> Result<i64> {
let now = Utc::now().timestamp();
let ts = if i.timestamp == 0 { now } else { i.timestamp };
let channel = if i.channel.is_empty() { "manual" } else { &i.channel };
store.conn().execute(
"INSERT INTO interactions (person_id, target_id, interaction_type,
channel, content, timestamp, created_at)
VALUES (?1,?2,?3,?4,?5,?6,?7)",
params![i.person_id, i.target_id, i.interaction_type,
channel, i.content, ts, now],
)?;
Ok(store.conn().last_insert_rowid())
}
pub fn interactions_for(store: &Store, person_id: i64) -> Result<Vec<Interaction>> {
let mut stmt = store.conn().prepare(
"SELECT id, person_id, target_id, interaction_type, channel, content, timestamp
FROM interactions WHERE person_id=?1 ORDER BY timestamp DESC",
)?;
let rows = stmt.query_map(params![person_id], |r| {
Ok(Interaction {
id: r.get(0)?, person_id: r.get(1)?, target_id: r.get(2)?,
interaction_type: r.get(3)?, channel: r.get(4)?,
content: r.get(5)?, timestamp: r.get(6)?,
})
})?;
let mut out = Vec::new();
for r in rows { out.push(r?); }
Ok(out)
}

View file

@ -0,0 +1,11 @@
//! kei-social-store — people + organizations + interactions.
pub mod graph;
pub mod interactions;
pub mod people;
pub mod schema;
pub mod search;
pub mod store;
pub use people::{Organization, Person};
pub use store::Store;

View file

@ -0,0 +1,80 @@
use clap::{Parser, Subcommand};
use kei_social_store::graph::relationship_graph;
use kei_social_store::interactions::{log_interaction, Interaction};
use kei_social_store::people::{add_org, add_person, Organization, Person};
use kei_social_store::search::search_people;
use kei_social_store::Store;
use std::path::PathBuf;
use std::process::ExitCode;
#[derive(Parser)]
#[command(name = "kei-social-store", version)]
struct Cli {
#[arg(long)] db: Option<PathBuf>,
#[command(subcommand)] cmd: Cmd,
}
#[derive(Subcommand)]
enum Cmd {
SearchPeople { query: String, #[arg(long, default_value_t = 20)] limit: i64 },
AddPerson { name: String,
#[arg(long, default_value = "")] email: String,
#[arg(long, default_value = "")] handle: String,
#[arg(long, default_value = "manual")] source: String },
AddOrg { name: String, #[arg(long, default_value = "company")] org_type: String },
LogInteraction { person_id: i64, interaction_type: String,
#[arg(long, default_value = "")] content: String,
#[arg(long, default_value = "manual")] channel: String,
#[arg(long, default_value_t = 0)] target_id: i64 },
RelationshipGraph,
}
fn db_path(o: Option<PathBuf>) -> PathBuf {
if let Some(p) = o { return p; }
if let Ok(e) = std::env::var("KEI_SOCIAL_DB") { return PathBuf::from(e); }
let home = std::env::var("HOME").unwrap_or_else(|_| ".".into());
PathBuf::from(home).join(".claude/social/social.sqlite")
}
fn run() -> anyhow::Result<()> {
let cli = Cli::parse();
let s = Store::open(&db_path(cli.db))?;
match cli.cmd {
Cmd::SearchPeople { query, limit } => {
for p in search_people(&s, &query, limit)? {
println!("{}\t{}\t{}", p.id, p.name, p.email);
}
}
Cmd::AddPerson { name, email, handle, source } => {
let id = add_person(&s, &Person {
name, email, handle, source, ..Default::default() })?;
println!("{}", id);
}
Cmd::AddOrg { name, org_type } => {
let id = add_org(&s, &Organization {
name, org_type, ..Default::default() })?;
println!("{}", id);
}
Cmd::LogInteraction { person_id, interaction_type, content, channel, target_id } => {
let id = log_interaction(&s, &Interaction {
person_id, target_id, interaction_type, channel, content,
..Default::default()
})?;
println!("{}", id);
}
Cmd::RelationshipGraph => {
for p in relationship_graph(&s)? {
println!("{}\t-[{}]->\t{}\t({}x)",
p.person_id, p.channel, p.target_id, p.count);
}
}
}
Ok(())
}
fn main() -> ExitCode {
match run() {
Ok(()) => ExitCode::SUCCESS,
Err(e) => { eprintln!("kei-social-store: {e:#}"); ExitCode::from(1) }
}
}

View file

@ -0,0 +1,76 @@
use crate::store::Store;
use anyhow::Result;
use chrono::Utc;
use rusqlite::params;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Person {
pub id: i64,
pub name: String,
pub email: String,
pub handle: String,
pub role: String,
pub organization: String,
pub source: String,
pub bio: String,
pub created_at: i64,
pub updated_at: i64,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Organization {
pub id: i64,
pub name: String,
pub org_type: String,
pub description: String,
pub created_at: i64,
}
pub fn add_person(store: &Store, p: &Person) -> Result<i64> {
let now = Utc::now().timestamp();
let source = if p.source.is_empty() { "manual" } else { &p.source };
store.conn().execute(
"INSERT INTO people (name, email, handle, role, organization,
source, bio, created_at, updated_at)
VALUES (?1,?2,?3,?4,?5,?6,?7,?8,?8)",
params![p.name, p.email, p.handle, p.role, p.organization,
source, p.bio, now],
)?;
let id = store.conn().last_insert_rowid();
store.conn().execute(
"INSERT INTO fts_social (person_id, name, email, bio) VALUES (?1,?2,?3,?4)",
params![id, p.name, p.email, p.bio],
)?;
Ok(id)
}
pub fn get_person(store: &Store, id: i64) -> Result<Option<Person>> {
let mut stmt = store.conn().prepare(
"SELECT id, name, email, handle, role, organization, source, bio,
created_at, updated_at FROM people WHERE id=?1",
)?;
let mut rows = stmt.query(params![id])?;
if let Some(r) = rows.next()? {
return Ok(Some(Person {
id: r.get(0)?, name: r.get(1)?, email: r.get(2)?, handle: r.get(3)?,
role: r.get(4)?, organization: r.get(5)?, source: r.get(6)?,
bio: r.get(7)?, created_at: r.get(8)?, updated_at: r.get(9)?,
}));
}
Ok(None)
}
pub fn add_org(store: &Store, o: &Organization) -> Result<i64> {
let now = Utc::now().timestamp();
let ot = if o.org_type.is_empty() { "company" } else { &o.org_type };
store.conn().execute(
"INSERT OR IGNORE INTO organizations (name, org_type, description, created_at)
VALUES (?1,?2,?3,?4)",
params![o.name, ot, o.description, now],
)?;
let id: i64 = store.conn().query_row(
"SELECT id FROM organizations WHERE name=?1",
params![o.name], |r| r.get(0))?;
Ok(id)
}

View file

@ -0,0 +1,48 @@
use rusqlite::{Connection, Result};
pub fn create_schema(conn: &Connection) -> Result<()> {
conn.execute_batch(r#"
CREATE TABLE IF NOT EXISTS people (
id INTEGER PRIMARY KEY,
name TEXT NOT NULL,
email TEXT DEFAULT '',
handle TEXT DEFAULT '',
role TEXT DEFAULT '',
organization TEXT DEFAULT '',
source TEXT NOT NULL DEFAULT 'manual',
bio TEXT DEFAULT '',
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL
);
CREATE UNIQUE INDEX IF NOT EXISTS idx_people_email
ON people(email) WHERE email != '';
CREATE UNIQUE INDEX IF NOT EXISTS idx_people_handle_source
ON people(handle, source) WHERE handle != '';
CREATE TABLE IF NOT EXISTS organizations (
id INTEGER PRIMARY KEY,
name TEXT NOT NULL UNIQUE,
org_type TEXT DEFAULT 'company',
description TEXT DEFAULT '',
created_at INTEGER NOT NULL
);
CREATE TABLE IF NOT EXISTS interactions (
id INTEGER PRIMARY KEY,
person_id INTEGER NOT NULL REFERENCES people(id) ON DELETE CASCADE,
target_id INTEGER NOT NULL DEFAULT 0,
interaction_type TEXT NOT NULL,
channel TEXT NOT NULL DEFAULT 'manual',
content TEXT DEFAULT '',
timestamp INTEGER NOT NULL,
created_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_int_person ON interactions(person_id);
"#)?;
conn.execute_batch(r#"
CREATE VIRTUAL TABLE IF NOT EXISTS fts_social
USING fts5(person_id UNINDEXED, name, email, bio,
tokenize='porter unicode61');
"#)?;
Ok(())
}

View file

@ -0,0 +1,25 @@
use crate::people::Person;
use crate::store::Store;
use anyhow::Result;
use rusqlite::params;
pub fn search_people(store: &Store, q: &str, limit: i64) -> Result<Vec<Person>> {
let lim = if limit <= 0 { 20 } else { limit };
let mut stmt = store.conn().prepare(
"SELECT p.id, p.name, p.email, p.handle, p.role, p.organization,
p.source, p.bio, p.created_at, p.updated_at
FROM fts_social f
JOIN people p ON p.id = f.person_id
WHERE fts_social MATCH ?1 ORDER BY rank LIMIT ?2",
)?;
let rows = stmt.query_map(params![q, lim], |r| {
Ok(Person {
id: r.get(0)?, name: r.get(1)?, email: r.get(2)?, handle: r.get(3)?,
role: r.get(4)?, organization: r.get(5)?, source: r.get(6)?,
bio: r.get(7)?, created_at: r.get(8)?, updated_at: r.get(9)?,
})
})?;
let mut out = Vec::new();
for r in rows { out.push(r?); }
Ok(out)
}

View file

@ -0,0 +1,22 @@
use crate::schema::create_schema;
use anyhow::{Context, Result};
use rusqlite::Connection;
use std::path::Path;
pub struct Store { conn: Connection }
impl Store {
pub fn open(path: &Path) -> Result<Self> {
if let Some(parent) = path.parent() { let _ = std::fs::create_dir_all(parent); }
let conn = Connection::open(path).context("open sqlite")?;
conn.pragma_update(None, "journal_mode", "WAL").ok();
create_schema(&conn)?;
Ok(Self { conn })
}
pub fn open_memory() -> Result<Self> {
let conn = Connection::open_in_memory()?;
create_schema(&conn)?;
Ok(Self { conn })
}
pub fn conn(&self) -> &Connection { &self.conn }
}

View file

@ -0,0 +1,68 @@
use kei_social_store::graph::relationship_graph;
use kei_social_store::interactions::{interactions_for, log_interaction, Interaction};
use kei_social_store::people::{add_org, add_person, get_person, Organization, Person};
use kei_social_store::search::search_people;
use kei_social_store::Store;
fn mk() -> Store { Store::open_memory().unwrap() }
#[test]
fn people_crud() {
let s = mk();
let id = add_person(&s, &Person {
name: "Alice".into(), email: "alice@example.com".into(),
..Default::default()
}).unwrap();
let p = get_person(&s, id).unwrap().unwrap();
assert_eq!(p.name, "Alice");
}
#[test]
fn orgs_idempotent() {
let s = mk();
let a = add_org(&s, &Organization { name: "Acme".into(), ..Default::default() }).unwrap();
let b = add_org(&s, &Organization { name: "Acme".into(), ..Default::default() }).unwrap();
assert_eq!(a, b);
}
#[test]
fn interactions_tracked() {
let s = mk();
let p = add_person(&s, &Person { name: "Bob".into(), ..Default::default() }).unwrap();
log_interaction(&s, &Interaction {
person_id: p, interaction_type: "email".into(),
content: "hi".into(), channel: "gmail".into(),
..Default::default()
}).unwrap();
let hist = interactions_for(&s, p).unwrap();
assert_eq!(hist.len(), 1);
assert_eq!(hist[0].interaction_type, "email");
}
#[test]
fn search_finds_person() {
let s = mk();
add_person(&s, &Person {
name: "Carol Chang".into(), bio: "rust async".into(),
..Default::default()
}).unwrap();
let hits = search_people(&s, "rust", 10).unwrap();
assert_eq!(hits.len(), 1);
assert_eq!(hits[0].name, "Carol Chang");
}
#[test]
fn relationship_graph_groups() {
let s = mk();
let a = add_person(&s, &Person { name: "A".into(), ..Default::default() }).unwrap();
let b = add_person(&s, &Person { name: "B".into(), ..Default::default() }).unwrap();
for _ in 0..3 {
log_interaction(&s, &Interaction {
person_id: a, target_id: b, interaction_type: "msg".into(),
channel: "slack".into(), ..Default::default()
}).unwrap();
}
let pairs = relationship_graph(&s).unwrap();
assert_eq!(pairs.len(), 1);
assert_eq!(pairs[0].count, 3);
}

View file

@ -0,0 +1,25 @@
[package]
name = "kei-task"
version = "0.1.0"
edition = "2021"
rust-version = "1.75"
description = "Task DAG with deps + milestones (SQLite). Port of LBM internal/task."
[[bin]]
name = "kei-task"
path = "src/main.rs"
[lib]
name = "kei_task"
path = "src/lib.rs"
[dependencies]
rusqlite = { version = "0.31", features = ["bundled"] }
clap = { version = "4", features = ["derive"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
anyhow = "1"
chrono = { version = "0.4", default-features = false, features = ["clock"] }
[dev-dependencies]
tempfile = "3"

View file

@ -0,0 +1,66 @@
//! Dependency edges + cycle detection + dependency-chain traversal.
use crate::store::Store;
use crate::types::is_valid_dep;
use anyhow::{anyhow, Result};
use rusqlite::params;
use std::collections::HashSet;
/// Add a dependency. Refuses a cycle (taskId -> dependsOn -> ... -> taskId).
pub fn add_dependency(store: &Store, task_id: i64, depends_on: i64, dep_type: &str) -> Result<()> {
if task_id == depends_on {
return Err(anyhow!("self-dependency forbidden"));
}
let dt = if dep_type.is_empty() { "blocks" } else { dep_type };
if !is_valid_dep(dt) {
return Err(anyhow!("invalid dep type: {dt}"));
}
if creates_cycle(store, task_id, depends_on)? {
return Err(anyhow!("cycle: {task_id} -> {depends_on} would close a loop"));
}
store.conn().execute(
"INSERT OR IGNORE INTO task_deps (task_id, depends_on, dep_type) VALUES (?1,?2,?3)",
params![task_id, depends_on, dt],
)?;
Ok(())
}
/// True if adding task_id -> depends_on would create a cycle.
fn creates_cycle(store: &Store, task_id: i64, depends_on: i64) -> Result<bool> {
// If depends_on reaches task_id via existing deps, cycle would close.
let mut stack = vec![depends_on];
let mut seen: HashSet<i64> = HashSet::new();
while let Some(cur) = stack.pop() {
if cur == task_id {
return Ok(true);
}
if !seen.insert(cur) {
continue;
}
let mut stmt = store.conn().prepare("SELECT depends_on FROM task_deps WHERE task_id=?1")?;
let rows = stmt.query_map(params![cur], |r| r.get::<_, i64>(0))?;
for row in rows {
stack.push(row?);
}
}
Ok(false)
}
/// Full dependency chain (BFS transitive closure).
pub fn dependency_chain(store: &Store, task_id: i64) -> Result<Vec<i64>> {
let mut seen: HashSet<i64> = HashSet::new();
let mut frontier = vec![task_id];
let mut chain: Vec<i64> = Vec::new();
while let Some(cur) = frontier.pop() {
let mut stmt = store.conn().prepare("SELECT depends_on FROM task_deps WHERE task_id=?1")?;
let rows = stmt.query_map(params![cur], |r| r.get::<_, i64>(0))?;
for row in rows {
let id = row?;
if seen.insert(id) {
chain.push(id);
frontier.push(id);
}
}
}
Ok(chain)
}

View file

@ -0,0 +1,28 @@
//! Adjacency view — returns task graph as edge-list for visualisation.
use crate::store::Store;
use anyhow::Result;
use serde::Serialize;
#[derive(Debug, Clone, Serialize)]
pub struct TaskEdge {
pub task_id: i64,
pub depends_on: i64,
pub dep_type: String,
}
pub fn list_edges(store: &Store) -> Result<Vec<TaskEdge>> {
let mut stmt = store.conn().prepare(
"SELECT task_id, depends_on, dep_type FROM task_deps"
)?;
let rows = stmt.query_map([], |r| {
Ok(TaskEdge {
task_id: r.get(0)?,
depends_on: r.get(1)?,
dep_type: r.get(2)?,
})
})?;
let mut out = Vec::new();
for r in rows { out.push(r?); }
Ok(out)
}

View file

@ -0,0 +1,12 @@
//! kei-task — tasks with typed deps (DAG, cycle-detected), milestones, FTS search.
pub mod deps;
pub mod graph;
pub mod milestones;
pub mod schema;
pub mod search;
pub mod store;
pub mod types;
pub use store::Store;
pub use types::{Milestone, Task};

View file

@ -0,0 +1,92 @@
//! kei-task CLI — create / update / add-dep / graph / dependency-chain.
use clap::{Parser, Subcommand};
use kei_task::deps::{add_dependency, dependency_chain};
use kei_task::graph::list_edges;
use kei_task::milestones::{create_milestone, link_task_to_milestone};
use kei_task::search::search;
use kei_task::{Milestone, Store, Task};
use std::path::PathBuf;
use std::process::ExitCode;
#[derive(Parser)]
#[command(name = "kei-task", version, about = "Task DAG CLI")]
struct Cli {
#[arg(long)] db: Option<PathBuf>,
#[command(subcommand)] cmd: Cmd,
}
#[derive(Subcommand)]
enum Cmd {
Create { title: String, #[arg(long, default_value = "")] description: String,
#[arg(long, default_value = "medium")] priority: String },
Update { id: i64, #[arg(long)] status: Option<String>, #[arg(long)] title: Option<String> },
AddDependency { from_id: i64, to_id: i64,
#[arg(long, default_value = "blocks")] dep_type: String },
Graph,
DependencyChain { id: i64 },
Search { query: String, #[arg(long, default_value_t = 20)] limit: i64 },
Milestone { name: String, #[arg(long, default_value = "")] description: String },
LinkMilestone { task_id: i64, milestone_id: i64 },
}
fn db_path(cli_db: Option<PathBuf>) -> PathBuf {
if let Some(p) = cli_db { return p; }
if let Ok(e) = std::env::var("KEI_TASK_DB") { return PathBuf::from(e); }
let home = std::env::var("HOME").unwrap_or_else(|_| ".".into());
PathBuf::from(home).join(".claude/task/task.sqlite")
}
fn run() -> anyhow::Result<()> {
let cli = Cli::parse();
let s = Store::open(&db_path(cli.db))?;
match cli.cmd {
Cmd::Create { title, description, priority } => {
let id = s.create_task(&Task { title, description, priority, ..Default::default() })?;
println!("{}", id);
}
Cmd::Update { id, status, title } => {
let mut t = s.get_task(id)?.ok_or_else(|| anyhow::anyhow!("id {id} not found"))?;
if let Some(st) = status { t.status = st; }
if let Some(ti) = title { t.title = ti; }
s.update_task(&t)?;
println!("updated {}", id);
}
Cmd::AddDependency { from_id, to_id, dep_type } => {
add_dependency(&s, from_id, to_id, &dep_type)?;
println!("dep: {} -> {} ({})", from_id, to_id, dep_type);
}
Cmd::Graph => {
for e in list_edges(&s)? {
println!("{}\t-[{}]->\t{}", e.task_id, e.dep_type, e.depends_on);
}
}
Cmd::DependencyChain { id } => {
for did in dependency_chain(&s, id)? {
println!("{}", did);
}
}
Cmd::Search { query, limit } => {
for t in search(&s, &query, limit)? {
println!("{}\t{}\t{}", t.id, t.status, t.title);
}
}
Cmd::Milestone { name, description } => {
let id = create_milestone(&s, &Milestone {
name, description, ..Default::default() })?;
println!("{}", id);
}
Cmd::LinkMilestone { task_id, milestone_id } => {
link_task_to_milestone(&s, task_id, milestone_id)?;
println!("linked {} -> milestone {}", task_id, milestone_id);
}
}
Ok(())
}
fn main() -> ExitCode {
match run() {
Ok(()) => ExitCode::SUCCESS,
Err(e) => { eprintln!("kei-task: {e:#}"); ExitCode::from(1) }
}
}

View file

@ -0,0 +1,39 @@
//! Milestone CRUD + task→milestone linking.
use crate::store::Store;
use crate::types::Milestone;
use anyhow::Result;
use chrono::Utc;
use rusqlite::params;
pub fn create_milestone(store: &Store, m: &Milestone) -> Result<i64> {
let now = Utc::now().timestamp();
let created = if m.created_at == 0 { now } else { m.created_at };
let status = if m.status.is_empty() { "open" } else { &m.status };
store.conn().execute(
"INSERT INTO milestones (name, description, target_date, status, created_at)
VALUES (?1,?2,?3,?4,?5)",
params![m.name, m.description, m.target_date, status, created],
)?;
Ok(store.conn().last_insert_rowid())
}
pub fn link_task_to_milestone(store: &Store, task_id: i64, milestone_id: i64) -> Result<()> {
store.conn().execute(
"INSERT OR IGNORE INTO task_milestones (task_id, milestone_id) VALUES (?1,?2)",
params![task_id, milestone_id],
)?;
Ok(())
}
pub fn tasks_in_milestone(store: &Store, milestone_id: i64) -> Result<Vec<i64>> {
let mut stmt = store.conn().prepare(
"SELECT task_id FROM task_milestones WHERE milestone_id=?1",
)?;
let rows = stmt.query_map(params![milestone_id], |r| r.get::<_, i64>(0))?;
let mut out = Vec::new();
for r in rows {
out.push(r?);
}
Ok(out)
}

View file

@ -0,0 +1,53 @@
//! SQLite schema for tasks + milestones + deps. Port of LBM internal/task/schema.go.
use rusqlite::{Connection, Result};
pub fn create_schema(conn: &Connection) -> Result<()> {
conn.execute_batch(r#"
CREATE TABLE IF NOT EXISTS tasks (
id INTEGER PRIMARY KEY,
title TEXT NOT NULL,
description TEXT DEFAULT '',
status TEXT NOT NULL DEFAULT 'pending',
priority TEXT DEFAULT 'medium',
task_type TEXT DEFAULT '',
parent_id INTEGER DEFAULT 0,
assigned_to TEXT DEFAULT '',
due_date INTEGER DEFAULT 0,
completed_at INTEGER DEFAULT 0,
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_task_status ON tasks(status);
CREATE INDEX IF NOT EXISTS idx_task_priority ON tasks(priority);
CREATE INDEX IF NOT EXISTS idx_task_parent ON tasks(parent_id);
CREATE TABLE IF NOT EXISTS milestones (
id INTEGER PRIMARY KEY,
name TEXT NOT NULL,
description TEXT DEFAULT '',
target_date INTEGER DEFAULT 0,
status TEXT DEFAULT 'open',
created_at INTEGER NOT NULL
);
CREATE TABLE IF NOT EXISTS task_deps (
task_id INTEGER NOT NULL,
depends_on INTEGER NOT NULL,
dep_type TEXT DEFAULT 'blocks',
PRIMARY KEY(task_id, depends_on)
);
CREATE INDEX IF NOT EXISTS idx_dep_depends ON task_deps(depends_on);
CREATE TABLE IF NOT EXISTS task_milestones (
task_id INTEGER NOT NULL,
milestone_id INTEGER NOT NULL,
PRIMARY KEY(task_id, milestone_id)
);
"#)?;
conn.execute_batch(r#"
CREATE VIRTUAL TABLE IF NOT EXISTS fts_tasks
USING fts5(task_id UNINDEXED, title, description, tokenize='porter unicode61');
"#)?;
Ok(())
}

View file

@ -0,0 +1,29 @@
//! FTS5 search over tasks (title + description).
use crate::store::Store;
use crate::types::Task;
use anyhow::Result;
use rusqlite::params;
pub fn search(store: &Store, query: &str, limit: i64) -> Result<Vec<Task>> {
let lim = if limit <= 0 { 20 } else { limit };
let mut stmt = store.conn().prepare(
"SELECT t.id, t.title, t.description, t.status, t.priority, t.task_type,
t.parent_id, t.assigned_to, t.due_date, t.completed_at,
t.created_at, t.updated_at
FROM fts_tasks f
JOIN tasks t ON t.id = f.task_id
WHERE fts_tasks MATCH ?1 ORDER BY rank LIMIT ?2",
)?;
let rows = stmt.query_map(params![query, lim], |r| {
Ok(Task {
id: r.get(0)?, title: r.get(1)?, description: r.get(2)?, status: r.get(3)?,
priority: r.get(4)?, task_type: r.get(5)?, parent_id: r.get(6)?,
assigned_to: r.get(7)?, due_date: r.get(8)?, completed_at: r.get(9)?,
created_at: r.get(10)?, updated_at: r.get(11)?,
})
})?;
let mut out = Vec::new();
for r in rows { out.push(r?); }
Ok(out)
}

View file

@ -0,0 +1,94 @@
//! Task store — open, CRUD, FTS reindex.
use crate::schema::create_schema;
use crate::types::Task;
use anyhow::{Context, Result};
use chrono::Utc;
use rusqlite::{params, Connection};
use std::path::Path;
pub struct Store {
conn: Connection,
}
impl Store {
pub fn open(path: &Path) -> Result<Self> {
if let Some(parent) = path.parent() {
let _ = std::fs::create_dir_all(parent);
}
let conn = Connection::open(path).context("open sqlite")?;
conn.pragma_update(None, "journal_mode", "WAL").ok();
create_schema(&conn)?;
Ok(Self { conn })
}
pub fn open_memory() -> Result<Self> {
let conn = Connection::open_in_memory()?;
create_schema(&conn)?;
Ok(Self { conn })
}
pub fn conn(&self) -> &Connection { &self.conn }
pub fn create_task(&self, t: &Task) -> Result<i64> {
let now = Utc::now().timestamp();
let created = if t.created_at == 0 { now } else { t.created_at };
let status = if t.status.is_empty() { "pending" } else { &t.status };
let priority = if t.priority.is_empty() { "medium" } else { &t.priority };
self.conn.execute(
"INSERT INTO tasks (title, description, status, priority, task_type,
parent_id, assigned_to, due_date, completed_at, created_at, updated_at)
VALUES (?1,?2,?3,?4,?5,?6,?7,?8,?9,?10,?11)",
params![t.title, t.description, status, priority, t.task_type,
t.parent_id, t.assigned_to, t.due_date, t.completed_at, created, now],
)?;
let id = self.conn.last_insert_rowid();
self.reindex_fts(id, &t.title, &t.description)?;
Ok(id)
}
pub fn get_task(&self, id: i64) -> Result<Option<Task>> {
let mut stmt = self.conn.prepare(
"SELECT id, title, description, status, priority, task_type, parent_id,
assigned_to, due_date, completed_at, created_at, updated_at
FROM tasks WHERE id=?1",
)?;
let mut rows = stmt.query(params![id])?;
if let Some(r) = rows.next()? {
return Ok(Some(row_to_task(r)?));
}
Ok(None)
}
pub fn update_task(&self, t: &Task) -> Result<()> {
let now = Utc::now().timestamp();
let completed = if t.status == "completed" && t.completed_at == 0 { now } else { t.completed_at };
self.conn.execute(
"UPDATE tasks SET title=?1, description=?2, status=?3, priority=?4,
task_type=?5, parent_id=?6, assigned_to=?7, due_date=?8,
completed_at=?9, updated_at=?10 WHERE id=?11",
params![t.title, t.description, t.status, t.priority, t.task_type,
t.parent_id, t.assigned_to, t.due_date, completed, now, t.id],
)?;
self.reindex_fts(t.id, &t.title, &t.description)?;
Ok(())
}
fn reindex_fts(&self, id: i64, title: &str, description: &str) -> Result<()> {
self.conn.execute("DELETE FROM fts_tasks WHERE task_id=?1", params![id])?;
self.conn.execute(
"INSERT INTO fts_tasks (task_id, title, description) VALUES (?1,?2,?3)",
params![id, title, description],
)?;
Ok(())
}
}
fn row_to_task(r: &rusqlite::Row) -> rusqlite::Result<Task> {
Ok(Task {
id: r.get(0)?, title: r.get(1)?, description: r.get(2)?, status: r.get(3)?,
priority: r.get(4)?, task_type: r.get(5)?, parent_id: r.get(6)?,
assigned_to: r.get(7)?, due_date: r.get(8)?, completed_at: r.get(9)?,
created_at: r.get(10)?, updated_at: r.get(11)?,
})
}

View file

@ -0,0 +1,45 @@
//! Task + Milestone value types and enum validation.
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Task {
pub id: i64,
pub title: String,
pub description: String,
pub status: String,
pub priority: String,
pub task_type: String,
pub parent_id: i64,
pub assigned_to: String,
pub due_date: i64,
pub completed_at: i64,
pub created_at: i64,
pub updated_at: i64,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Milestone {
pub id: i64,
pub name: String,
pub description: String,
pub target_date: i64,
pub status: String,
pub created_at: i64,
}
pub const VALID_STATUSES: &[&str] =
&["pending", "in_progress", "completed", "cancelled", "blocked"];
pub const VALID_PRIORITIES: &[&str] = &["critical", "high", "medium", "low"];
pub const VALID_DEP_TYPES: &[&str] =
&["blocks", "feeds_into", "subtask_of", "milestone_of", "assigned_to", "depends_on"];
pub fn is_valid_status(s: &str) -> bool {
VALID_STATUSES.contains(&s)
}
pub fn is_valid_priority(s: &str) -> bool {
VALID_PRIORITIES.contains(&s)
}
pub fn is_valid_dep(s: &str) -> bool {
VALID_DEP_TYPES.contains(&s)
}

View file

@ -0,0 +1,95 @@
//! kei-task integration tests.
use kei_task::deps::{add_dependency, dependency_chain};
use kei_task::graph::list_edges;
use kei_task::milestones::{create_milestone, link_task_to_milestone, tasks_in_milestone};
use kei_task::search::search;
use kei_task::{Milestone, Store, Task};
fn mk() -> Store { Store::open_memory().unwrap() }
fn mktask(title: &str) -> Task {
Task { title: title.into(), priority: "high".into(), ..Default::default() }
}
#[test]
fn create_and_get() {
let s = mk();
let id = s.create_task(&mktask("a")).unwrap();
let t = s.get_task(id).unwrap().unwrap();
assert_eq!(t.title, "a");
assert_eq!(t.status, "pending");
}
#[test]
fn update_persists() {
let s = mk();
let id = s.create_task(&mktask("a")).unwrap();
let mut t = s.get_task(id).unwrap().unwrap();
t.status = "in_progress".into();
s.update_task(&t).unwrap();
let u = s.get_task(id).unwrap().unwrap();
assert_eq!(u.status, "in_progress");
}
#[test]
fn cycle_detected() {
let s = mk();
let a = s.create_task(&mktask("a")).unwrap();
let b = s.create_task(&mktask("b")).unwrap();
let c = s.create_task(&mktask("c")).unwrap();
add_dependency(&s, a, b, "blocks").unwrap();
add_dependency(&s, b, c, "blocks").unwrap();
// a -> b -> c; now c -> a would be a cycle
let err = add_dependency(&s, c, a, "blocks");
assert!(err.is_err(), "cycle detection must reject");
}
#[test]
fn milestone_linking() {
let s = mk();
let t = s.create_task(&mktask("design")).unwrap();
let ms_id = create_milestone(&s, &Milestone {
name: "v1".into(), ..Default::default() }).unwrap();
link_task_to_milestone(&s, t, ms_id).unwrap();
let tasks = tasks_in_milestone(&s, ms_id).unwrap();
assert_eq!(tasks, vec![t]);
}
#[test]
fn dependency_chain_traversal() {
let s = mk();
let a = s.create_task(&mktask("a")).unwrap();
let b = s.create_task(&mktask("b")).unwrap();
let c = s.create_task(&mktask("c")).unwrap();
add_dependency(&s, a, b, "blocks").unwrap();
add_dependency(&s, b, c, "blocks").unwrap();
let chain = dependency_chain(&s, a).unwrap();
assert!(chain.contains(&b));
assert!(chain.contains(&c));
assert_eq!(chain.len(), 2);
}
#[test]
fn task_graph_edges() {
let s = mk();
let a = s.create_task(&mktask("a")).unwrap();
let b = s.create_task(&mktask("b")).unwrap();
add_dependency(&s, a, b, "blocks").unwrap();
let edges = list_edges(&s).unwrap();
assert_eq!(edges.len(), 1);
assert_eq!(edges[0].task_id, a);
assert_eq!(edges[0].depends_on, b);
}
#[test]
fn search_finds_task() {
let s = mk();
s.create_task(&Task {
title: "refactor router".into(),
description: "split monolith".into(),
..Default::default()
}).unwrap();
let hits = search(&s, "refactor", 10).unwrap();
assert_eq!(hits.len(), 1);
}