checkpoint: matrix-agent-validated (2026-04-25)

Architectural snapshot of the lakehouse codebase at the point where the
full matrix-driven agent loop with Mem0 versioning + deletion was
validated end-to-end.

WHAT THIS REPO IS
A clean single-commit snapshot of the lakehouse code. Heavy test data
(.parquet datasets, vector indexes) excluded — see REPLICATION.md for
regen path. Full lakehouse history at git.agentview.dev/profit/lakehouse.

WHAT WAS PROVEN
- Vector retrieval across multi-corpora matrix (chicago_permits + entity
  briefs + sec_tickers + distilled procedural + llm_team runs)
- Observer hand-review (cloud + heuristic fallback) gating each candidate
- Local-model agent loop (qwen3.5:latest) with tool use + scratchpad
- Playbook seal on success → next-iter retrieval surfaces it as preamble
- Mem0 versioning + deletion in pathway_memory:
    * UPSERT: ADD on new workflow, UPDATE bumps replay_count on identical
    * REVISE: chains versions, parent.superseded_at + superseded_by stamped
    * RETIRE: marks specific trace retired with reason, excluded from retrieval
    * HISTORY: walks chain root→tip, cycle-safe

KEY DIRECTORIES
- crates/vectord/src/pathway_memory.rs — Mem0 ops live here
- crates/vectord/src/playbook_memory.rs — original Mem0 reference
- tests/agent_test/ — local-model agent harness + PRD + session archives
- scripts/dump_raw_corpus.sh — MinIO bucket dump (raw test corpus)
- scripts/vectorize_raw_corpus.ts — corpus → vector indexes
- scripts/analyze_chicago_contracts.ts — real inference pipeline
- scripts/seal_agent_playbook.ts — Mem0 upsert from agent traces

Replication: see REPLICATION.md for Debian 13 clean install + cloud-only
adaptation (no local Ollama).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
profit
2026-04-25 19:43:27 -05:00
co-authored by Claude Opus 4.7
commit ac01fffd9a
325 changed files with 82720 additions and 0 deletions
+14
View File
@@ -0,0 +1,14 @@
[package]
name = "aibridge"
version = "0.1.0"
edition = "2024"
[dependencies]
shared = { path = "../shared" }
tokio = { workspace = true }
axum = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }
tracing = { workspace = true }
reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] }
async-trait = "0.1"
+188
View File
@@ -0,0 +1,188 @@
use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::time::Duration;
/// HTTP client for the Python AI sidecar.
#[derive(Clone)]
pub struct AiClient {
client: Client,
base_url: String,
}
// -- Request/Response types --
#[derive(Serialize, Deserialize)]
pub struct EmbedRequest {
pub texts: Vec<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub model: Option<String>,
}
#[derive(Deserialize, Serialize, Clone)]
pub struct EmbedResponse {
pub embeddings: Vec<Vec<f64>>,
pub model: String,
pub dimensions: usize,
}
#[derive(Clone, Serialize, Deserialize)]
pub struct GenerateRequest {
pub prompt: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub model: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub system: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub temperature: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_tokens: Option<u32>,
/// Phase 21 — per-call opt-out of hidden reasoning. Thinking models
/// (qwen3.5, gpt-oss, etc) burn tokens on reasoning before the
/// visible response starts; setting this to `false` on hot-path
/// JSON emitters avoids empty returns when the budget is tight.
/// Sidecar forwards this to Ollama's `think` parameter; if the
/// sidecar drops an unknown field the request still succeeds.
#[serde(skip_serializing_if = "Option::is_none")]
pub think: Option<bool>,
}
#[derive(Deserialize, Serialize, Clone)]
pub struct GenerateResponse {
pub text: String,
pub model: String,
pub tokens_evaluated: Option<u64>,
pub tokens_generated: Option<u64>,
}
#[derive(Serialize, Deserialize)]
pub struct RerankRequest {
pub query: String,
pub documents: Vec<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub model: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub top_k: Option<usize>,
}
#[derive(Deserialize, Serialize, Clone)]
pub struct ScoredDocument {
pub index: usize,
pub text: String,
pub score: f64,
}
#[derive(Deserialize, Serialize, Clone)]
pub struct RerankResponse {
pub results: Vec<ScoredDocument>,
pub model: String,
}
impl AiClient {
pub fn new(base_url: &str) -> Self {
let client = Client::builder()
.timeout(Duration::from_secs(120))
.build()
.expect("failed to build HTTP client");
Self {
client,
base_url: base_url.trim_end_matches('/').to_string(),
}
}
pub async fn health(&self) -> Result<serde_json::Value, String> {
let resp = self.client
.get(format!("{}/health", self.base_url))
.send()
.await
.map_err(|e| format!("sidecar unreachable: {e}"))?;
resp.json().await.map_err(|e| format!("invalid response: {e}"))
}
pub async fn embed(&self, req: EmbedRequest) -> Result<EmbedResponse, String> {
let resp = self.client
.post(format!("{}/embed", self.base_url))
.json(&req)
.send()
.await
.map_err(|e| format!("embed request failed: {e}"))?;
if !resp.status().is_success() {
let text = resp.text().await.unwrap_or_default();
return Err(format!("embed error ({}): {text}", text.len()));
}
resp.json().await.map_err(|e| format!("embed parse error: {e}"))
}
pub async fn generate(&self, req: GenerateRequest) -> Result<GenerateResponse, String> {
let resp = self.client
.post(format!("{}/generate", self.base_url))
.json(&req)
.send()
.await
.map_err(|e| format!("generate request failed: {e}"))?;
if !resp.status().is_success() {
let text = resp.text().await.unwrap_or_default();
return Err(format!("generate error: {text}"));
}
resp.json().await.map_err(|e| format!("generate parse error: {e}"))
}
pub async fn rerank(&self, req: RerankRequest) -> Result<RerankResponse, String> {
let resp = self.client
.post(format!("{}/rerank", self.base_url))
.json(&req)
.send()
.await
.map_err(|e| format!("rerank request failed: {e}"))?;
if !resp.status().is_success() {
let text = resp.text().await.unwrap_or_default();
return Err(format!("rerank error: {text}"));
}
resp.json().await.map_err(|e| format!("rerank parse error: {e}"))
}
/// Force Ollama to unload the named model from VRAM (keep_alive=0).
/// Used for predictable profile swaps — without this, Ollama holds a
/// model for its configured TTL (default 5min) and the previous
/// profile's model can linger in VRAM next to the new one.
pub async fn unload_model(&self, model: &str) -> Result<serde_json::Value, String> {
let resp = self.client
.post(format!("{}/admin/unload", self.base_url))
.json(&serde_json::json!({ "model": model }))
.send().await
.map_err(|e| format!("unload request failed: {e}"))?;
if !resp.status().is_success() {
let text = resp.text().await.unwrap_or_default();
return Err(format!("unload error: {text}"));
}
resp.json().await.map_err(|e| format!("unload parse error: {e}"))
}
/// Ask Ollama to load the named model into VRAM proactively. Makes
/// the first real request after profile activation fast (no cold-load
/// latency).
pub async fn preload_model(&self, model: &str) -> Result<serde_json::Value, String> {
let resp = self.client
.post(format!("{}/admin/preload", self.base_url))
.json(&serde_json::json!({ "model": model }))
.send().await
.map_err(|e| format!("preload request failed: {e}"))?;
if !resp.status().is_success() {
let text = resp.text().await.unwrap_or_default();
return Err(format!("preload error: {text}"));
}
resp.json().await.map_err(|e| format!("preload parse error: {e}"))
}
/// GPU + loaded-model snapshot from the sidecar. Combines nvidia-smi
/// output (if available) with Ollama's /api/ps.
pub async fn vram_snapshot(&self) -> Result<serde_json::Value, String> {
let resp = self.client
.get(format!("{}/admin/vram", self.base_url))
.send().await
.map_err(|e| format!("vram request failed: {e}"))?;
resp.json().await.map_err(|e| format!("vram parse error: {e}"))
}
}
+188
View File
@@ -0,0 +1,188 @@
//! Phase 21 — context-budget accounting for model calls.
//!
//! Ports `assertContextBudget` + `estimateTokens` + `CONTEXT_WINDOWS`
//! from `tests/multi-agent/agent.ts` so Rust-side callers (gateway
//! tool surfaces, future Rust agents) get the same loud-fail behavior
//! on window overflow instead of silent truncation.
//!
//! The token estimator is deliberately the same chars/4 heuristic as
//! the TS side. It's biased ~15% safe — pessimistic on English, correct
//! within a factor of 2 on code. Swap to a provider tokenizer only when
//! the estimator drives a decision (we're nowhere near that yet).
use std::collections::HashMap;
use std::sync::OnceLock;
// `estimate_tokens` moved to `shared::model_matrix::ModelMatrix::estimate_tokens`
// (cdc24d8). All callers migrated; the deprecated wrapper that stood in its
// place has been removed since it had zero external consumers.
/// Phase 21 — per-model context windows, mirroring the TS table in
/// `tests/multi-agent/agent.ts`. Anchored on each model's documented
/// max; unknown models fall back to `DEFAULT_CONTEXT_WINDOW`.
pub const DEFAULT_CONTEXT_WINDOW: usize = 32_768;
pub const DEFAULT_SAFETY_MARGIN: usize = 2_000;
pub const DEFAULT_MAX_TOKENS: usize = 800;
fn known_windows() -> &'static HashMap<&'static str, usize> {
static TABLE: OnceLock<HashMap<&'static str, usize>> = OnceLock::new();
TABLE.get_or_init(|| {
let mut m = HashMap::new();
m.insert("mistral:latest", 32_768);
m.insert("qwen2.5:latest", 32_768);
m.insert("qwen3:latest", 40_960);
m.insert("qwen3.5:latest", 262_144);
m.insert("qwen3-embedding", 32_768);
m.insert("nomic-embed-text-v2-moe", 2_048);
m.insert("gpt-oss:20b", 131_072);
m.insert("gpt-oss:120b", 131_072);
m.insert("qwen3.5:397b", 131_072);
m.insert("kimi-k2-thinking", 200_000);
m.insert("kimi-k2.6", 200_000);
m.insert("kimi-k2:1t", 1_048_576);
m.insert("deepseek-v3.1:671b", 131_072);
m.insert("glm-4.7", 131_072);
m
})
}
pub fn context_window_for(model: &str) -> usize {
known_windows().get(model).copied().unwrap_or(DEFAULT_CONTEXT_WINDOW)
}
/// Result of a budget check — exposes the numbers so callers can log
/// how much headroom remains without re-running the estimator.
#[derive(Debug, Clone, Copy)]
pub struct BudgetCheck {
pub estimated: usize,
pub window: usize,
pub remaining: i64,
}
/// Inputs to `assert_context_budget`. `bypass` exists for call sites
/// that handle their own overflow (continuation's second pass already
/// counted the partial; T5 gatekeeper prompts have a separate policy).
#[derive(Debug, Clone, Default)]
pub struct BudgetOpts<'a> {
pub system: Option<&'a str>,
pub max_tokens: Option<usize>,
pub safety_margin: Option<usize>,
pub bypass: bool,
}
/// Phase 21's loud-fail primitive. Returns a `BudgetCheck` on success
/// and the same struct plus over-by count on failure. The whole point
/// is to stop silent truncation — callers that expect overflow should
/// chunk BEFORE calling or set `bypass: true`.
pub fn assert_context_budget(
model: &str,
prompt: &str,
opts: BudgetOpts,
) -> Result<BudgetCheck, (BudgetCheck, usize)> {
let window = context_window_for(model);
let safety = opts.safety_margin.unwrap_or(DEFAULT_SAFETY_MARGIN);
let max_tokens = opts.max_tokens.unwrap_or(DEFAULT_MAX_TOKENS);
let sys_tokens = opts.system.map(shared::model_matrix::ModelMatrix::estimate_tokens).unwrap_or(0);
let estimated = shared::model_matrix::ModelMatrix::estimate_tokens(prompt) + sys_tokens + max_tokens;
let remaining = window as i64 - estimated as i64 - safety as i64;
let check = BudgetCheck { estimated, window, remaining };
if remaining < 0 && !opts.bypass {
return Err((check, (-remaining) as usize));
}
Ok(check)
}
/// Convenience — format an overflow error the same way the TS side
/// does. Exposed so downstream crates render consistent messages.
pub fn overflow_message(model: &str, check: &BudgetCheck, over_by: usize, safety: usize) -> String {
format!(
"context overflow: model={} est={}t window={}t safety={}t over={}t. \
Chunk the prompt (see config/models.json overflow_policies) or set \
bypass:true if you know the risk.",
model, check.estimated, check.window, safety, over_by,
)
}
#[cfg(test)]
mod tests {
use super::*;
// Deprecated-function behavior is now canonically tested in
// crates/shared/src/model_matrix.rs. This test was the legacy
// pin that preceded the migration; delete when the deprecated
// wrapper itself goes (see the #[deprecated] attribute).
#[test]
fn context_window_known_and_fallback() {
assert_eq!(context_window_for("qwen3.5:latest"), 262_144);
assert_eq!(context_window_for("kimi-k2:1t"), 1_048_576);
assert_eq!(context_window_for("some-unreleased-model"), DEFAULT_CONTEXT_WINDOW);
}
#[test]
fn budget_passes_well_under_window() {
let check = assert_context_budget(
"qwen3:latest",
&"x".repeat(4_000), // ~1000 tokens
BudgetOpts { max_tokens: Some(500), ..Default::default() },
).expect("well under 40K window");
assert!(check.remaining > 30_000);
}
#[test]
fn budget_fails_when_prompt_overflows_window() {
let huge = "x".repeat(200_000); // ~50K tokens, over qwen3's 40K
let err = assert_context_budget(
"qwen3:latest",
&huge,
BudgetOpts::default(),
).expect_err("should overflow qwen3's 40K window");
assert!(err.1 > 0, "over_by must be positive");
}
#[test]
fn budget_bypass_returns_ok_even_over() {
let huge = "x".repeat(200_000);
let check = assert_context_budget(
"qwen3:latest",
&huge,
BudgetOpts { bypass: true, ..Default::default() },
).expect("bypass must suppress the error");
assert!(check.remaining < 0, "check still reports negative remaining");
}
#[test]
fn budget_counts_system_prompt() {
// 10K-char system prompt → ~2500 tokens. With a big max_tokens
// this should push us closer to the window.
let sys = "s".repeat(10_000);
let prompt = "p".repeat(4_000);
let with_sys = assert_context_budget(
"qwen3:latest",
&prompt,
BudgetOpts {
system: Some(&sys),
max_tokens: Some(500),
..Default::default()
},
).unwrap();
let without_sys = assert_context_budget(
"qwen3:latest",
&prompt,
BudgetOpts { max_tokens: Some(500), ..Default::default() },
).unwrap();
assert!(with_sys.estimated > without_sys.estimated,
"system prompt should raise estimate");
assert_eq!(with_sys.estimated - without_sys.estimated, shared::model_matrix::ModelMatrix::estimate_tokens(&sys));
}
#[test]
fn overflow_message_includes_numbers() {
let check = BudgetCheck { estimated: 42_000, window: 40_960, remaining: -1_040 };
let msg = overflow_message("qwen3:latest", &check, 3_040, 2_000);
assert!(msg.contains("qwen3:latest"));
assert!(msg.contains("42000t"));
assert!(msg.contains("40960t"));
assert!(msg.contains("3040t"));
}
}
+461
View File
@@ -0,0 +1,461 @@
//! Phase 21 — OUTPUT-overflow handler. Ports `generateContinuable`
//! from `tests/multi-agent/agent.ts`.
//!
//! Two failure modes to repair:
//!
//! * EMPTY response — thinking model ate the entire budget on hidden
//! reasoning before emitting a token. Fix: retry the original prompt
//! with 2× the budget, geometric up to `BUDGET_CAP`.
//!
//! * TRUNCATED non-empty — model got most of the way but hit
//! max_tokens before closing the structure. Fix: continue with the
//! partial response in the prompt as scratchpad, so the model knows
//! where to pick up without restarting.
//!
//! `TextGenerator` abstracts the sidecar so tests can inject canned
//! responses without a live Ollama.
use std::future::Future;
use crate::client::{AiClient, GenerateRequest, GenerateResponse};
/// Shape classifier for `is_structurally_complete`. JSON responses
/// must parse; text responses just need to be non-empty.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ResponseShape {
Json,
Text,
}
/// Trait that `generate_continuable` + `generate_tree_split` call. The
/// real implementation forwards to `AiClient::generate`; tests supply a
/// mock with a scripted sequence of responses.
pub trait TextGenerator: Send + Sync {
fn generate_text(
&self,
req: GenerateRequest,
) -> impl Future<Output = Result<GenerateResponse, String>> + Send;
}
impl TextGenerator for AiClient {
fn generate_text(
&self,
req: GenerateRequest,
) -> impl Future<Output = Result<GenerateResponse, String>> + Send {
self.generate(req)
}
}
/// Strip a surrounding ```json``` fence if present. Leaves the inner
/// content alone otherwise. Returns a slice of `s`.
fn strip_json_fence(s: &str) -> &str {
let t = s.trim();
if let Some(rest) = t.strip_prefix("```json") {
rest.trim_start_matches('\n').strip_suffix("```").unwrap_or(rest).trim()
} else if let Some(rest) = t.strip_prefix("```") {
rest.trim_start_matches('\n').strip_suffix("```").unwrap_or(rest).trim()
} else {
t
}
}
/// Port of the TS brace-balance + JSON.parse check. Returns true when
/// the outermost `{...}` block is balanced and parses. Text shape is
/// satisfied by any non-empty, non-whitespace payload.
pub fn is_structurally_complete(text: &str, shape: ResponseShape) -> bool {
if text.trim().is_empty() { return false; }
if shape == ResponseShape::Text { return true; }
let s = strip_json_fence(text);
let Some(start) = s.find('{') else { return false };
let Some(end) = s.rfind('}') else { return false };
if end <= start { return false; }
let slice = &s[start..=end];
// Balance check — cheaper than parse, catches truncated nests.
// String state tracked because `{` inside a string doesn't count.
let mut depth: i32 = 0;
let mut in_str = false;
let mut esc = false;
for c in slice.chars() {
if esc { esc = false; continue; }
if c == '\\' { esc = true; continue; }
if c == '"' { in_str = !in_str; continue; }
if in_str { continue; }
if c == '{' { depth += 1; }
else if c == '}' {
depth -= 1;
if depth < 0 { return false; }
}
}
if depth != 0 { return false; }
// Parse check is the tie-breaker — balanced but invalid JSON (e.g.
// trailing comma before `}`) shouldn't count as complete.
serde_json::from_str::<serde_json::Value>(slice).is_ok()
}
/// Knobs for `generate_continuable`. All optional with sensible
/// defaults that match the TS version.
#[derive(Debug, Clone)]
pub struct ContinuableOpts {
pub model: String,
pub max_tokens: Option<u32>,
pub temperature: Option<f64>,
pub system: Option<String>,
pub shape: ResponseShape,
pub max_continuations: usize,
pub think: Option<bool>,
/// Geometric-backoff ceiling for the empty-response retry path.
/// Matches TS's `budgetCap = 8000`.
pub budget_cap: u32,
/// Maximum empty-response retries before giving up. Matches TS's
/// hardcoded `retry < 3`.
pub max_empty_retries: usize,
}
impl ContinuableOpts {
pub fn new(model: impl Into<String>) -> Self {
Self {
model: model.into(),
max_tokens: None,
temperature: None,
system: None,
shape: ResponseShape::Json,
max_continuations: 3,
think: None,
budget_cap: 8_000,
max_empty_retries: 3,
}
}
}
/// Outcome of a `generate_continuable` call. Carries the combined
/// text plus diagnostic counters so observability downstream can
/// report "how many continuations did that query cost".
#[derive(Debug, Clone)]
pub struct ContinuableOutcome {
pub text: String,
pub empty_retries: usize,
pub continuations: usize,
pub final_complete: bool,
/// Sum of `prompt_tokens` across every generator call made to
/// produce this outcome — including empty retries and continuations.
/// Lets callers (gateway execution loop, observability) stamp
/// accurate per-task usage without second-guessing the retry fan-out.
pub prompt_tokens: u32,
/// Sum of `completion_tokens` across every generator call.
pub completion_tokens: u32,
/// Total number of generator calls. `1 + empty_retries +
/// continuations` in the normal case; the field is explicit so
/// callers don't have to re-derive it.
pub calls: u32,
}
fn make_request(opts: &ContinuableOpts, prompt: String, current_max: u32) -> GenerateRequest {
GenerateRequest {
prompt,
model: Some(opts.model.clone()),
system: opts.system.clone(),
temperature: opts.temperature,
max_tokens: Some(current_max),
think: opts.think,
}
}
fn continuation_prompt(original: &str, partial: &str) -> String {
format!(
"{original}\n\n\
PARTIAL RESPONSE SO FAR (continue from here — do NOT restart, \
do NOT repeat what's already there, emit ONLY the remaining \
tokens to close the structure):\n{partial}"
)
}
/// Phase 21 — output-overflow safe generate. See module docs for the
/// two failure modes repaired. On final-still-incomplete, returns the
/// combined text with `final_complete: false` so the caller's parser
/// can throw with the raw text for forensics rather than silently
/// truncating.
pub async fn generate_continuable<G: TextGenerator>(
generator: &G,
prompt: &str,
opts: &ContinuableOpts,
) -> Result<ContinuableOutcome, String> {
let initial_max = opts.max_tokens.unwrap_or(800);
let mut current_max = initial_max;
let mut combined = String::new();
let mut empty_retries = 0usize;
let mut continuations = 0usize;
let mut prompt_tokens: u32 = 0;
let mut completion_tokens: u32 = 0;
let mut calls: u32 = 0;
// Phase 21(a) — empty-response backoff loop.
for retry in 0..opts.max_empty_retries {
let req = make_request(opts, prompt.to_string(), current_max);
let resp = generator.generate_text(req).await?;
calls += 1;
prompt_tokens = prompt_tokens.saturating_add(resp.tokens_evaluated.unwrap_or(0) as u32);
completion_tokens = completion_tokens.saturating_add(resp.tokens_generated.unwrap_or(0) as u32);
if !resp.text.trim().is_empty() {
combined = resp.text;
break;
}
empty_retries = retry + 1;
current_max = (current_max.saturating_mul(2)).min(opts.budget_cap);
}
// Phase 21(b) — structural-completion continuation loop.
for _ in 0..opts.max_continuations {
if is_structurally_complete(&combined, opts.shape) {
return Ok(ContinuableOutcome {
text: combined,
empty_retries,
continuations,
final_complete: true,
prompt_tokens,
completion_tokens,
calls,
});
}
if combined.trim().is_empty() {
// Nothing to continue from — continuing "" is identical to
// the initial call and would loop.
break;
}
let cont_prompt = continuation_prompt(prompt, &combined);
let req = make_request(opts, cont_prompt, current_max.min(opts.budget_cap));
let resp = generator.generate_text(req).await?;
calls += 1;
prompt_tokens = prompt_tokens.saturating_add(resp.tokens_evaluated.unwrap_or(0) as u32);
completion_tokens = completion_tokens.saturating_add(resp.tokens_generated.unwrap_or(0) as u32);
combined.push_str(&resp.text);
continuations += 1;
}
let final_complete = is_structurally_complete(&combined, opts.shape);
Ok(ContinuableOutcome {
text: combined,
empty_retries,
continuations,
final_complete,
prompt_tokens,
completion_tokens,
calls,
})
}
/// Scripted generator for unit tests. Returns responses from `script`
/// in order; extra calls reuse the last entry so tests don't have to
/// count past what they actually assert on.
#[cfg(test)]
pub struct ScriptedGenerator {
script: Vec<Result<String, String>>,
calls: std::sync::Arc<std::sync::Mutex<Vec<GenerateRequest>>>,
}
#[cfg(test)]
impl ScriptedGenerator {
pub fn new<I>(script: I) -> Self
where
I: IntoIterator<Item = Result<String, String>>,
{
Self {
script: script.into_iter().collect(),
calls: std::sync::Arc::new(std::sync::Mutex::new(Vec::new())),
}
}
pub fn calls(&self) -> Vec<GenerateRequest> {
self.calls.lock().unwrap().clone()
}
pub fn call_count(&self) -> usize {
self.calls.lock().unwrap().len()
}
}
#[cfg(test)]
impl TextGenerator for ScriptedGenerator {
fn generate_text(
&self,
req: GenerateRequest,
) -> impl Future<Output = Result<GenerateResponse, String>> + Send {
let i = {
let mut calls = self.calls.lock().unwrap();
let i = calls.len();
calls.push(req.clone());
i
};
let model = req.model.clone().unwrap_or_default();
let entry = self.script.get(i)
.cloned()
.unwrap_or_else(|| self.script.last().cloned().unwrap_or(Ok(String::new())));
async move {
entry.map(|text| GenerateResponse {
text,
model,
tokens_evaluated: None,
tokens_generated: None,
})
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn structural_complete_rejects_empty_and_text_mismatch() {
assert!(!is_structurally_complete("", ResponseShape::Text));
assert!(!is_structurally_complete(" ", ResponseShape::Text));
assert!(is_structurally_complete("any content", ResponseShape::Text));
}
#[test]
fn structural_complete_handles_balanced_json() {
assert!(is_structurally_complete(r#"{"a": 1}"#, ResponseShape::Json));
assert!(is_structurally_complete(
r#"```json
{"a": 1, "b": [1, 2, 3]}
```"#,
ResponseShape::Json,
));
}
#[test]
fn structural_complete_rejects_truncated_json() {
assert!(!is_structurally_complete(r#"{"a": 1"#, ResponseShape::Json));
assert!(!is_structurally_complete(r#"{"a": {"b": 1"#, ResponseShape::Json));
// Trailing comma — balanced but unparseable
assert!(!is_structurally_complete(r#"{"a": 1,}"#, ResponseShape::Json));
}
#[test]
fn structural_complete_ignores_braces_inside_strings() {
assert!(is_structurally_complete(r#"{"s": "has } inside"}"#, ResponseShape::Json));
assert!(is_structurally_complete(r#"{"s": "escaped \" quote"}"#, ResponseShape::Json));
}
#[tokio::test]
async fn continuable_returns_first_response_when_complete() {
let generator = ScriptedGenerator::new(vec![Ok(r#"{"ok": true}"#.to_string())]);
let opts = ContinuableOpts::new("qwen3:latest");
let out = generate_continuable(&generator, "test", &opts).await.unwrap();
assert!(out.final_complete);
assert_eq!(out.empty_retries, 0);
assert_eq!(out.continuations, 0);
assert_eq!(generator.call_count(), 1);
}
#[tokio::test]
async fn continuable_retries_on_empty_with_doubled_budget() {
// Two empties, then a good response. Third call should see 4×
// the initial budget (2× twice).
let generator = ScriptedGenerator::new(vec![
Ok("".to_string()),
Ok("".to_string()),
Ok(r#"{"ok": true}"#.to_string()),
]);
let mut opts = ContinuableOpts::new("qwen3:latest");
opts.max_tokens = Some(100);
let out = generate_continuable(&generator, "test", &opts).await.unwrap();
assert!(out.final_complete);
assert_eq!(out.empty_retries, 2);
let calls = generator.calls();
assert_eq!(calls.len(), 3);
assert_eq!(calls[0].max_tokens, Some(100));
assert_eq!(calls[1].max_tokens, Some(200));
assert_eq!(calls[2].max_tokens, Some(400));
}
#[tokio::test]
async fn continuable_caps_budget_at_budget_cap() {
let generator = ScriptedGenerator::new(vec![
Ok("".to_string()),
Ok("".to_string()),
Ok(r#"{"ok":1}"#.to_string()),
]);
let mut opts = ContinuableOpts::new("qwen3:latest");
opts.max_tokens = Some(5_000);
opts.budget_cap = 7_000;
let out = generate_continuable(&generator, "test", &opts).await.unwrap();
assert!(out.final_complete);
let calls = generator.calls();
// 5000 → doubled would be 10000; cap pulls it to 7000.
assert_eq!(calls[1].max_tokens, Some(7_000));
assert_eq!(calls[2].max_tokens, Some(7_000));
}
#[tokio::test]
async fn continuable_glues_truncated_response() {
// First call returns balanced-open `{...`; continuation closes
// it with `...}`. Combined must parse.
let generator = ScriptedGenerator::new(vec![
Ok(r#"{"fills": [{"candidate_id": "C-001""#.to_string()),
Ok(r#", "name": "Alice"}]}"#.to_string()),
]);
let opts = ContinuableOpts::new("qwen3:latest");
let out = generate_continuable(&generator, "ORIGINAL", &opts).await.unwrap();
assert!(out.final_complete, "combined must parse: {}", out.text);
assert_eq!(out.continuations, 1);
let calls = generator.calls();
assert_eq!(calls.len(), 2);
// Continuation prompt must contain the partial — that's the
// scratchpad primitive J called out.
let cont_prompt = &calls[1].prompt;
assert!(cont_prompt.contains("ORIGINAL"),
"continuation must include original prompt");
assert!(cont_prompt.contains(r#"{"fills": [{"candidate_id": "C-001""#),
"continuation must include partial");
}
#[tokio::test]
async fn continuable_does_not_loop_on_persistent_empty() {
// All three retries return empty; we must NOT then enter the
// continuation loop with an empty partial (would burn 3 more
// calls continuing from "").
let generator = ScriptedGenerator::new(vec![
Ok("".to_string()),
Ok("".to_string()),
Ok("".to_string()),
]);
let opts = ContinuableOpts::new("qwen3:latest");
let out = generate_continuable(&generator, "test", &opts).await.unwrap();
assert!(!out.final_complete);
assert_eq!(out.empty_retries, 3);
assert_eq!(out.continuations, 0, "must not continue from empty");
assert_eq!(generator.call_count(), 3);
}
#[tokio::test]
async fn continuable_returns_raw_on_exhausted_continuations() {
// Three continuations that never complete — caller's parser
// will throw. We must return the combined text so forensics
// has the raw content, not a lossy truncation.
let generator = ScriptedGenerator::new(vec![
Ok(r#"{"a": ["#.to_string()),
Ok("1,".to_string()),
Ok("2,".to_string()),
Ok("3,".to_string()),
]);
let mut opts = ContinuableOpts::new("qwen3:latest");
opts.max_continuations = 3;
let out = generate_continuable(&generator, "test", &opts).await.unwrap();
assert!(!out.final_complete);
assert_eq!(out.continuations, 3);
assert!(out.text.contains(r#"{"a": ["#));
assert!(out.text.contains("3,"));
}
#[tokio::test]
async fn continuable_propagates_generator_errors() {
let generator = ScriptedGenerator::new(vec![Err("sidecar 503".to_string())]);
let opts = ContinuableOpts::new("qwen3:latest");
let err = generate_continuable(&generator, "test", &opts).await.unwrap_err();
assert!(err.contains("503"));
}
}
+8
View File
@@ -0,0 +1,8 @@
pub mod client;
pub mod context;
pub mod continuation;
pub mod provider;
pub mod providers;
pub mod routing;
pub mod service;
pub mod tree_split;
+39
View File
@@ -0,0 +1,39 @@
use async_trait::async_trait;
use crate::client::{GenerateRequest, GenerateResponse, EmbedRequest, EmbedResponse};
#[async_trait]
pub trait ProviderAdapter: Send + Sync {
/// Name for routing (ollama, openrouter, etc.)
fn name(&self) -> &str;
/// Chat completion — returns text, model, token counts
async fn chat(&self, req: GenerateRequest) -> Result<GenerateResponse, String>;
/// Embeddings — returns vectors, model, dimensions
async fn embed(&self, req: EmbedRequest) -> Result<EmbedResponse, String>;
/// Unload model from VRAM (optional, no-op if not supported)
async fn unload(&self, _model: &str) -> Result<(), String> {
Ok(())
}
/// Health check
async fn health(&self) -> Result<serde_json::Value, String>;
}
/// Routing key extracted from model name.
/// - "qwen3.5:latest" → "ollama"
/// - "openrouter/anthropic/claude-3.5-sonnet" → "openrouter"
/// - "gpt-4o" → "ollama" (default)
pub fn provider_key(model: &str) -> &'static str {
let lower = model.to_lowercase();
if lower.starts_with("openrouter/") {
"openrouter"
} else if lower.starts_with("gemini") {
"gemini"
} else if lower.starts_with("claude") {
"claude"
} else {
"ollama" // default: local Ollama
}
}
+2
View File
@@ -0,0 +1,2 @@
pub mod ollama;
pub mod openrouter;
+37
View File
@@ -0,0 +1,37 @@
use async_trait::async_trait;
use crate::client::{AiClient, GenerateRequest, GenerateResponse, EmbedRequest, EmbedResponse};
use crate::provider::ProviderAdapter;
pub struct OllamaAdapter {
client: AiClient,
}
impl OllamaAdapter {
pub fn new(client: AiClient) -> Self {
Self { client }
}
}
#[async_trait]
impl ProviderAdapter for OllamaAdapter {
fn name(&self) -> &str {
"ollama"
}
async fn chat(&self, req: GenerateRequest) -> Result<GenerateResponse, String> {
self.client.generate(req).await
}
async fn embed(&self, req: EmbedRequest) -> Result<EmbedResponse, String> {
self.client.embed(req).await
}
async fn unload(&self, model: &str) -> Result<(), String> {
let _: serde_json::Value = self.client.unload_model(model).await?;
Ok(())
}
async fn health(&self) -> Result<serde_json::Value, String> {
self.client.health().await
}
}
+152
View File
@@ -0,0 +1,152 @@
use async_trait::async_trait;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::time::Duration;
use crate::client::{GenerateRequest, GenerateResponse, EmbedRequest, EmbedResponse};
use crate::provider::ProviderAdapter;
pub struct OpenRouterAdapter {
client: Client,
base_url: String,
api_key: String,
default_model: String,
}
#[derive(Serialize)]
struct OpenRouterChatRequest {
model: String,
messages: Vec<OpenRouterMessage>,
temperature: Option<f64>,
max_tokens: Option<u32>,
}
#[derive(Serialize)]
struct OpenRouterMessage {
role: String,
content: String,
}
#[derive(Deserialize)]
struct OpenRouterChatResponse {
choices: Vec<OpenRouterChoice>,
usage: OpenRouterUsage,
model: String,
}
#[derive(Deserialize)]
struct OpenRouterChoice {
message: OpenRouterMessageOut,
}
#[derive(Deserialize)]
#[allow(dead_code)]
struct OpenRouterMessageOut {
role: String,
content: String,
}
#[derive(Deserialize)]
#[allow(dead_code)]
struct OpenRouterUsage {
prompt_tokens: Option<u32>,
completion_tokens: Option<u32>,
total_tokens: Option<u32>,
}
impl OpenRouterAdapter {
pub fn new(base_url: &str, api_key: String, default_model: &str) -> Self {
let client = Client::builder()
.timeout(Duration::from_secs(180))
.build()
.expect("failed to build HTTP client");
Self {
client,
base_url: base_url.trim_end_matches('/').to_string(),
api_key,
default_model: default_model.to_string(),
}
}
fn chat_model(&self, model: &str) -> String {
// Strip "openrouter/" prefix if present
let m = model.trim_start_matches("openrouter/");
if m.is_empty() || m == model {
self.default_model.clone()
} else {
m.to_string()
}
}
fn to_openrouter_messages(req: &GenerateRequest) -> Vec<OpenRouterMessage> {
let mut out = vec![];
if let Some(sys) = &req.system {
out.push(OpenRouterMessage { role: "system".into(), content: sys.clone() });
}
out.push(OpenRouterMessage {
role: "user".into(),
content: req.prompt.clone(),
});
out
}
}
#[async_trait]
impl ProviderAdapter for OpenRouterAdapter {
fn name(&self) -> &str {
"openrouter"
}
async fn chat(&self, req: GenerateRequest) -> Result<GenerateResponse, String> {
let model = self.chat_model(req.model.as_deref().unwrap_or(""));
let or_req = OpenRouterChatRequest {
model: model.clone(),
messages: OpenRouterAdapter::to_openrouter_messages(&req),
temperature: req.temperature,
max_tokens: req.max_tokens,
};
let resp = self.client
.post(format!("{}/chat/completions", self.base_url))
.header("Authorization", format!("Bearer {}", self.api_key))
.header("Content-Type", "application/json")
.json(&or_req)
.send()
.await
.map_err(|e| format!("openrouter request failed: {e}"))?;
let status = resp.status();
let body = resp.text().await.unwrap_or_default();
if !status.is_success() {
return Err(format!("openrouter error ({}): {}", status, body));
}
let or_resp: OpenRouterChatResponse = serde_json::from_str(&body)
.map_err(|e| format!("openrouter parse error: {e}"))?;
let choice = or_resp.choices.into_iter().next()
.ok_or("no completion choice returned")?;
let usage = or_resp.usage;
Ok(GenerateResponse {
text: choice.message.content,
model: or_resp.model,
tokens_evaluated: usage.prompt_tokens.map(|n| n as u64),
tokens_generated: usage.completion_tokens.map(|n| n as u64),
})
}
async fn embed(&self, _req: EmbedRequest) -> Result<EmbedResponse, String> {
Err("openrouter: embed not implemented".into())
}
async fn health(&self) -> Result<serde_json::Value, String> {
// OpenRouter doesn't have a dedicated health endpoint,
// so we just return a healthy response if the client works.
Ok(serde_json::json!({
"status": "ok",
"provider": "openrouter",
}))
}
}
+119
View File
@@ -0,0 +1,119 @@
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct RoutingRule {
pub model_pattern: String,
pub provider: String,
pub max_tokens: Option<u32>,
pub temperature: Option<f64>,
}
#[derive(Clone, Default)]
pub struct RoutingEngine {
rules: Vec<RoutingRule>,
fallback_chain: Vec<String>,
}
impl RoutingEngine {
pub fn new() -> Self {
Self::default()
}
pub fn with_rules(mut self, rules: Vec<RoutingRule>) -> Self {
self.rules = rules;
self
}
pub fn with_fallback(mut self, chain: Vec<String>) -> Self {
self.fallback_chain = chain;
self
}
pub fn route(&self, model: &str) -> RouteDecision {
let lower = model.to_lowercase();
for rule in &self.rules {
if glob_match(&rule.model_pattern.to_lowercase(), &lower) {
return RouteDecision {
provider: rule.provider.clone(),
model: model.to_string(),
max_tokens: rule.max_tokens,
temperature: rule.temperature,
};
}
}
if let Some(first) = self.fallback_chain.first() {
RouteDecision {
provider: first.clone(),
model: model.to_string(),
max_tokens: None,
temperature: None,
}
} else {
RouteDecision {
provider: "ollama".to_string(),
model: model.to_string(),
max_tokens: None,
temperature: None,
}
}
}
}
#[derive(Clone, Debug)]
pub struct RouteDecision {
pub provider: String,
pub model: String,
pub max_tokens: Option<u32>,
pub temperature: Option<f64>,
}
fn glob_match(pattern: &str, name: &str) -> bool {
if !pattern.contains('*') { return pattern == name; }
let parts: Vec<&str> = pattern.split('*').collect();
// Multi-* support: first must be prefix, last must be suffix, each
// interior piece must appear in order. Fixes the iter-9 finding
// where gpt-*-large* silently fell through to an exact-match path.
// Also removes the dead `parts.len() == 1` branch that accessed
// parts[1] and would panic if ever reached (unreachable today
// since split('*') on a string containing '*' always yields ≥2).
if !name.starts_with(parts[0]) || !name.ends_with(parts.last().unwrap()) { return false; }
let mut cursor = parts[0].len();
parts[1..parts.len() - 1].iter().all(|mid| {
name[cursor..].find(mid).map(|pos| { cursor += pos + mid.len(); true }).unwrap_or(false)
})
}
impl Default for RoutingRule {
fn default() -> Self {
Self {
model_pattern: "*".to_string(),
provider: "ollama".to_string(),
max_tokens: None,
temperature: None,
}
}
}
#[cfg(test)]
mod glob_match_tests {
use super::glob_match;
#[test] fn exact_match() { assert!(glob_match("gpt-oss:120b", "gpt-oss:120b")); }
#[test] fn exact_mismatch() { assert!(!glob_match("a", "b")); }
#[test] fn leading_wildcard() { assert!(glob_match("*:120b", "gpt-oss:120b")); }
#[test] fn trailing_wildcard() { assert!(glob_match("gpt-oss:*", "gpt-oss:120b")); }
#[test] fn bare_wildcard() { assert!(glob_match("*", "anything")); }
#[test] fn multi_wildcard_in_order() { assert!(glob_match("gpt-*-oss-*", "gpt-4-oss-120b")); }
#[test] fn multi_wildcard_wrong_order() { assert!(!glob_match("b*a*", "abba")); }
#[test] fn multi_wildcard_panic_safety() {
// Regression: earlier impl had an unreachable `parts.len() == 1`
// branch that indexed parts[1] — would panic if ever hit. Now
// the split('*') invariant guarantees ≥2 parts when * present,
// and we handle all N-part cases explicitly.
assert!(glob_match("a*b*c", "abc"));
assert!(glob_match("a*b*c", "axxxbxxxc"));
assert!(!glob_match("a*b*c", "xxxbxxx"));
}
}
+68
View File
@@ -0,0 +1,68 @@
use axum::{
Json, Router,
extract::State,
http::StatusCode,
response::IntoResponse,
routing::{get, post},
};
use crate::client::{
AiClient, EmbedRequest, GenerateRequest, RerankRequest,
};
pub fn router(client: AiClient) -> Router {
Router::new()
.route("/health", get(health))
.route("/embed", post(embed))
.route("/generate", post(generate))
.route("/rerank", post(rerank))
.route("/vram", get(vram))
.with_state(client)
}
async fn vram(State(client): State<AiClient>) -> impl IntoResponse {
match client.vram_snapshot().await {
Ok(snap) => Ok(Json(snap)),
Err(e) => Err((StatusCode::BAD_GATEWAY, e)),
}
}
async fn health(State(client): State<AiClient>) -> impl IntoResponse {
match client.health().await {
Ok(info) => Ok(Json(info)),
Err(e) => Err((StatusCode::BAD_GATEWAY, format!("sidecar down: {e}"))),
}
}
async fn embed(
State(client): State<AiClient>,
Json(req): Json<EmbedRequest>,
) -> impl IntoResponse {
tracing::info!("embedding {} texts", req.texts.len());
match client.embed(req).await {
Ok(resp) => Ok(Json(resp)),
Err(e) => Err((StatusCode::BAD_GATEWAY, e)),
}
}
async fn generate(
State(client): State<AiClient>,
Json(req): Json<GenerateRequest>,
) -> impl IntoResponse {
tracing::info!("generating with prompt len={}", req.prompt.len());
match client.generate(req).await {
Ok(resp) => Ok(Json(resp)),
Err(e) => Err((StatusCode::BAD_GATEWAY, e)),
}
}
async fn rerank(
State(client): State<AiClient>,
Json(req): Json<RerankRequest>,
) -> impl IntoResponse {
tracing::info!("reranking {} documents", req.documents.len());
match client.rerank(req).await {
Ok(resp) => Ok(Json(resp)),
Err(e) => Err((StatusCode::BAD_GATEWAY, e)),
}
}
+327
View File
@@ -0,0 +1,327 @@
//! Phase 21 — INPUT-overflow handler. Ports `generateTreeSplit` from
//! `tests/multi-agent/agent.ts`.
//!
//! When the input corpus exceeds the model's window (200 playbooks
//! pasted into a T4 strategic prompt, a long retrospective digest, a
//! cross-corpus summarization), raising `max_tokens` doesn't help —
//! the prompt itself is the problem. The answer is map-reduce:
//!
//! 1. Caller shards the input at semantic boundaries (records,
//! paragraphs, playbook entries).
//! 2. For each shard, build a map prompt that includes the running
//! scratchpad and run it through `generate_continuable`.
//! 3. Append the map output to the scratchpad (oldest-first
//! truncation when it outgrows `scratchpad_budget`).
//! 4. Build a reduce prompt from the final scratchpad and run it.
//!
//! Every shard prompt and the reduce prompt go through
//! `assert_context_budget` first — if a single shard still overflows
//! we bubble the error up rather than silently truncating. That's the
//! whole point of Phase 21.
use crate::context::{assert_context_budget, BudgetOpts, overflow_message,
DEFAULT_MAX_TOKENS, DEFAULT_SAFETY_MARGIN};
use shared::model_matrix::ModelMatrix;
use crate::continuation::{generate_continuable, ContinuableOpts, ResponseShape, TextGenerator};
/// Callback signatures — caller supplies closures that stitch the
/// scratchpad into each shard's prompt and build the final reduce
/// prompt. Kept as `Fn` (not `FnMut`) so the map loop can call them
/// by reference.
pub type MapPromptFn<'a> = dyn Fn(&str, &str) -> String + Send + Sync + 'a;
pub type ReducePromptFn<'a> = dyn Fn(&str) -> String + Send + Sync + 'a;
/// Knobs for `generate_tree_split`.
#[derive(Debug, Clone)]
pub struct TreeSplitOpts {
pub model: String,
pub system: Option<String>,
pub temperature: Option<f64>,
/// max_tokens for map AND reduce (reduce defaults are usually
/// higher; caller overrides for just reduce by calling through
/// continuable directly if needed).
pub max_tokens: Option<u32>,
pub reduce_max_tokens: Option<u32>,
pub think: Option<bool>,
/// Soft ceiling on scratchpad size (estimated tokens). When it
/// grows past this, the oldest shard digest gets dropped. Default
/// 6000, matching the TS implementation.
pub scratchpad_budget: usize,
pub safety_margin: Option<usize>,
}
impl TreeSplitOpts {
pub fn new(model: impl Into<String>) -> Self {
Self {
model: model.into(),
system: None,
temperature: None,
max_tokens: None,
reduce_max_tokens: None,
think: None,
scratchpad_budget: 6_000,
safety_margin: None,
}
}
}
/// Result — final reduce response plus the accumulated scratchpad so
/// the caller can inspect what was kept vs truncated.
#[derive(Debug, Clone)]
pub struct TreeSplitResult {
pub response: String,
pub scratchpad: String,
pub shards_processed: usize,
pub scratchpad_truncations: usize,
pub total_continuations: usize,
}
/// Drop shard-digest blocks from the head of `scratchpad` until its
/// estimated-token count fits the budget. Digest blocks are delimited
/// by `\n— shard N/M digest —\n` so we can find the first one and
/// chop everything before its successor.
fn truncate_scratchpad(scratchpad: &mut String, budget_tokens: usize) -> bool {
if ModelMatrix::estimate_tokens(scratchpad) <= budget_tokens { return false; }
// Find the second delimiter — everything before it gets dropped.
const DELIM_PREFIX: &str = "\n— shard ";
let mut cursor = 0;
let mut truncated = false;
while ModelMatrix::estimate_tokens(&scratchpad[cursor..]) > budget_tokens {
// Skip past a leading delimiter (if we're sitting on one from
// a previous iteration), then find the next.
let search_from = cursor + if scratchpad[cursor..].starts_with(DELIM_PREFIX) {
DELIM_PREFIX.len()
} else { 0 };
let Some(rel_next) = scratchpad[search_from..].find(DELIM_PREFIX) else { break };
cursor = search_from + rel_next;
truncated = true;
}
if cursor > 0 {
scratchpad.drain(..cursor);
}
truncated
}
/// Phase 21 — map-reduce over shards with a running scratchpad. See
/// module docs.
pub async fn generate_tree_split<G: TextGenerator>(
generator: &G,
shards: &[String],
map_prompt: &MapPromptFn<'_>,
reduce_prompt: &ReducePromptFn<'_>,
opts: &TreeSplitOpts,
) -> Result<TreeSplitResult, String> {
let mut scratchpad = String::new();
let safety = opts.safety_margin.unwrap_or(DEFAULT_SAFETY_MARGIN);
let map_max = opts.max_tokens.unwrap_or(DEFAULT_MAX_TOKENS as u32);
let reduce_max = opts.reduce_max_tokens.unwrap_or(1_500);
let mut truncations = 0usize;
let mut total_continuations = 0usize;
for (i, shard) in shards.iter().enumerate() {
let shard_prompt = map_prompt(shard, &scratchpad);
// Loud-fail on per-shard overflow — caller sharded too
// coarsely. Silent truncation is exactly the mode J rejected.
let budget = BudgetOpts {
system: opts.system.as_deref(),
max_tokens: Some(map_max as usize),
safety_margin: Some(safety),
bypass: false,
};
let check = assert_context_budget(&opts.model, &shard_prompt, budget)
.map_err(|(c, over)| overflow_message(&opts.model, &c, over, safety))?;
let _ = check;
let mut cont_opts = ContinuableOpts::new(&opts.model);
cont_opts.max_tokens = Some(map_max);
cont_opts.temperature = opts.temperature;
cont_opts.system = opts.system.clone();
cont_opts.shape = ResponseShape::Text;
cont_opts.think = opts.think;
let outcome = generate_continuable(generator, &shard_prompt, &cont_opts).await?;
total_continuations += outcome.continuations;
// Append this shard's digest and, if needed, drop oldest.
scratchpad.push_str(&format!(
"\n— shard {}/{} digest —\n{}",
i + 1, shards.len(), outcome.text.trim(),
));
if truncate_scratchpad(&mut scratchpad, opts.scratchpad_budget) {
truncations += 1;
}
}
// Reduce pass. Budget check first — if the scratchpad is still too
// big for the reduce prompt we fail loud with numbers.
let reduce_p = reduce_prompt(&scratchpad);
let budget = BudgetOpts {
system: opts.system.as_deref(),
max_tokens: Some(reduce_max as usize),
safety_margin: Some(safety),
bypass: false,
};
assert_context_budget(&opts.model, &reduce_p, budget)
.map_err(|(c, over)| overflow_message(&opts.model, &c, over, safety))?;
let mut cont_opts = ContinuableOpts::new(&opts.model);
cont_opts.max_tokens = Some(reduce_max);
cont_opts.temperature = opts.temperature;
cont_opts.system = opts.system.clone();
cont_opts.shape = ResponseShape::Text;
cont_opts.think = opts.think;
let outcome = generate_continuable(generator, &reduce_p, &cont_opts).await?;
total_continuations += outcome.continuations;
Ok(TreeSplitResult {
response: outcome.text,
scratchpad,
shards_processed: shards.len(),
scratchpad_truncations: truncations,
total_continuations,
})
}
#[cfg(test)]
mod tests {
use super::*;
use crate::continuation::ScriptedGenerator;
fn simple_map(shard: &str, scratchpad: &str) -> String {
format!("SCRATCHPAD:\n{scratchpad}\n---\nSHARD:\n{shard}\n---\nDIGEST:")
}
fn simple_reduce(scratchpad: &str) -> String {
format!("SCRATCHPAD:\n{scratchpad}\n---\nFINAL:")
}
#[tokio::test]
async fn tree_split_runs_map_then_reduce() {
// 3 shards → 3 map calls → 1 reduce call = 4 responses scripted.
let generator = ScriptedGenerator::new(vec![
Ok("digest-1".to_string()),
Ok("digest-2".to_string()),
Ok("digest-3".to_string()),
Ok("FINAL ANSWER".to_string()),
]);
let shards: Vec<String> = ["a", "b", "c"].iter().map(|s| s.to_string()).collect();
let opts = TreeSplitOpts::new("qwen3:latest");
let map_fn: &MapPromptFn = &simple_map;
let reduce_fn: &ReducePromptFn = &simple_reduce;
let result = generate_tree_split(&generator, &shards, map_fn, reduce_fn, &opts)
.await
.unwrap();
assert_eq!(result.shards_processed, 3);
assert_eq!(result.response, "FINAL ANSWER");
assert_eq!(generator.call_count(), 4);
// Scratchpad must carry all three digests in order.
assert!(result.scratchpad.contains("digest-1"));
assert!(result.scratchpad.contains("digest-2"));
assert!(result.scratchpad.contains("digest-3"));
}
#[tokio::test]
async fn tree_split_reduce_prompt_sees_full_scratchpad() {
let generator = ScriptedGenerator::new(vec![
Ok("summary-A".to_string()),
Ok("summary-B".to_string()),
Ok("REDUCED".to_string()),
]);
let shards = vec!["input-one".to_string(), "input-two".to_string()];
let opts = TreeSplitOpts::new("qwen3:latest");
let _ = generate_tree_split(&generator, &shards, &simple_map, &simple_reduce, &opts)
.await
.unwrap();
// Third call = reduce. Its prompt must include both digests.
let calls = generator.calls();
let reduce_prompt = &calls[2].prompt;
assert!(reduce_prompt.contains("summary-A"),
"reduce prompt must see first shard digest");
assert!(reduce_prompt.contains("summary-B"),
"reduce prompt must see second shard digest");
}
#[tokio::test]
async fn tree_split_loud_fails_on_shard_overflow() {
let generator = ScriptedGenerator::new(vec![Ok("digest".to_string())]);
// One gigantic shard — well over qwen3's 40K window even as a
// prompt. The budget check must reject before any generate call.
let shards = vec!["x".repeat(200_000)];
let opts = TreeSplitOpts::new("qwen3:latest");
let err = generate_tree_split(&generator, &shards, &simple_map, &simple_reduce, &opts)
.await
.expect_err("shard-sized overflow must be rejected");
assert!(err.contains("overflow"), "error should mention overflow: {err}");
assert_eq!(generator.call_count(), 0, "generate must not be called on overflow");
}
#[tokio::test]
async fn tree_split_truncates_scratchpad_when_over_budget() {
// Tight budget so each shard trips truncation. qwen3's 40K
// window is fine; the budget we care about is the scratchpad
// cap, not the model window.
let generator = ScriptedGenerator::new(vec![
Ok("A".repeat(2_000)),
Ok("B".repeat(2_000)),
Ok("C".repeat(2_000)),
Ok("D".repeat(2_000)),
Ok("FINAL".to_string()),
]);
let shards: Vec<String> = (0..4).map(|i| format!("shard{i}")).collect();
let mut opts = TreeSplitOpts::new("qwen3:latest");
opts.scratchpad_budget = 1_000; // ~4000 chars — one digest barely fits
let result = generate_tree_split(&generator, &shards, &simple_map, &simple_reduce, &opts)
.await
.unwrap();
assert!(result.scratchpad_truncations > 0,
"tight budget must trigger truncation");
// Scratchpad should still fit roughly within the budget
// (post-truncation); the estimator uses chars/4 so the bound
// is ~budget*4 chars. Give some slack for the delimiter.
let scratchpad_tokens = ModelMatrix::estimate_tokens(&result.scratchpad);
assert!(scratchpad_tokens <= opts.scratchpad_budget * 2,
"scratchpad {} tokens vs budget {}", scratchpad_tokens, opts.scratchpad_budget);
}
#[tokio::test]
async fn tree_split_reports_continuations_from_map_and_reduce() {
// First shard: truncated-then-continued. Reduce: truncated-then-continued.
// 1 shard: 2 map calls (initial + continuation), then 2 reduce calls.
let generator = ScriptedGenerator::new(vec![
Ok("partial".to_string()), // map shape=text, non-empty → complete on first pass
Ok("reduce-out".to_string()),
]);
let shards = vec!["only".to_string()];
let opts = TreeSplitOpts::new("qwen3:latest");
let result = generate_tree_split(&generator, &shards, &simple_map, &simple_reduce, &opts)
.await
.unwrap();
// Text shape treats non-empty as complete → 0 continuations.
assert_eq!(result.total_continuations, 0);
assert_eq!(result.shards_processed, 1);
}
#[test]
fn truncate_scratchpad_noop_when_under_budget() {
let mut s = "\n— shard 1/1 digest —\nshort".to_string();
let truncated = truncate_scratchpad(&mut s, 1_000);
assert!(!truncated);
assert!(s.contains("short"));
}
#[test]
fn truncate_scratchpad_drops_oldest_first() {
let mut s = format!(
"\n— shard 1/3 digest —\n{}\n— shard 2/3 digest —\n{}\n— shard 3/3 digest —\nshort",
"x".repeat(4_000), // ~1000 tokens
"y".repeat(4_000), // ~1000 tokens
);
let truncated = truncate_scratchpad(&mut s, 500); // ~2000 chars
assert!(truncated);
assert!(!s.contains(&"x".repeat(4_000)),
"oldest digest should be dropped");
assert!(s.contains("short"),
"newest digest should survive");
}
}
+20
View File
@@ -0,0 +1,20 @@
[package]
name = "catalogd"
version = "0.1.0"
edition = "2024"
[dependencies]
shared = { path = "../shared" }
storaged = { path = "../storaged" }
tokio = { workspace = true }
axum = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }
tracing = { workspace = true }
bytes = { workspace = true }
chrono = { workspace = true }
uuid = { workspace = true }
object_store = { workspace = true }
arrow = { workspace = true }
proto = { path = "../proto" }
tonic = { workspace = true }
+104
View File
@@ -0,0 +1,104 @@
use proto::lakehouse::{
catalog_service_server::CatalogService,
CreateDatasetRequest, DatasetResponse, GetDatasetByNameRequest,
GetDatasetRequest, ListDatasetsRequest, ListDatasetsResponse,
ObjectRef as ProtoObjectRef,
};
use shared::types::{DatasetId, ObjectRef, SchemaFingerprint};
use tonic::{Request, Response, Status};
use uuid::Uuid;
use crate::registry::Registry;
pub struct CatalogGrpc {
registry: Registry,
}
impl CatalogGrpc {
pub fn new(registry: Registry) -> Self {
Self { registry }
}
}
fn manifest_to_proto(m: &shared::types::DatasetManifest) -> DatasetResponse {
DatasetResponse {
id: m.id.to_string(),
name: m.name.clone(),
schema_fingerprint: m.schema_fingerprint.0.clone(),
objects: m.objects.iter().map(|o| ProtoObjectRef {
bucket: o.bucket.clone(),
key: o.key.clone(),
size_bytes: o.size_bytes,
created_at: o.created_at.to_rfc3339(),
}).collect(),
created_at: m.created_at.to_rfc3339(),
updated_at: m.updated_at.to_rfc3339(),
}
}
#[tonic::async_trait]
impl CatalogService for CatalogGrpc {
async fn create_dataset(
&self,
request: Request<CreateDatasetRequest>,
) -> Result<Response<DatasetResponse>, Status> {
let req = request.into_inner();
let now = chrono::Utc::now();
let objects: Vec<ObjectRef> = req.objects.into_iter().map(|o| ObjectRef {
bucket: o.bucket,
key: o.key,
size_bytes: o.size_bytes,
created_at: now,
}).collect();
let manifest = self.registry
.register(req.name, SchemaFingerprint(req.schema_fingerprint), objects)
.await
.map_err(|e| {
// Mirror the HTTP side's 409 mapping: schema drift is a
// precondition failure, not a server error. Keeps gRPC
// and HTTP callers seeing the same diagnostic signal.
if e.contains("different schema") {
Status::failed_precondition(e)
} else {
Status::internal(e)
}
})?;
Ok(Response::new(manifest_to_proto(&manifest)))
}
async fn get_dataset(
&self,
request: Request<GetDatasetRequest>,
) -> Result<Response<DatasetResponse>, Status> {
let req = request.into_inner();
let uuid = Uuid::parse_str(&req.id).map_err(|e| Status::invalid_argument(e.to_string()))?;
let id = DatasetId(uuid);
match self.registry.get(&id).await {
Some(m) => Ok(Response::new(manifest_to_proto(&m))),
None => Err(Status::not_found(format!("dataset not found: {}", req.id))),
}
}
async fn get_dataset_by_name(
&self,
request: Request<GetDatasetByNameRequest>,
) -> Result<Response<DatasetResponse>, Status> {
let req = request.into_inner();
match self.registry.get_by_name(&req.name).await {
Some(m) => Ok(Response::new(manifest_to_proto(&m))),
None => Err(Status::not_found(format!("dataset not found: {}", req.name))),
}
}
async fn list_datasets(
&self,
_request: Request<ListDatasetsRequest>,
) -> Result<Response<ListDatasetsResponse>, Status> {
let datasets = self.registry.list().await;
let responses = datasets.iter().map(manifest_to_proto).collect();
Ok(Response::new(ListDatasetsResponse { datasets: responses }))
}
}
+4
View File
@@ -0,0 +1,4 @@
pub mod registry;
pub mod service;
pub mod grpc;
pub mod tombstones;
File diff suppressed because it is too large Load Diff
+475
View File
@@ -0,0 +1,475 @@
use axum::{
Json, Router,
extract::{Path, State},
http::StatusCode,
response::IntoResponse,
routing::{get, post},
};
use serde::{Deserialize, Serialize};
use shared::types::{DatasetId, ObjectRef, SchemaFingerprint};
use uuid::Uuid;
use crate::registry::Registry;
pub fn router(registry: Registry) -> Router {
Router::new()
.route("/health", get(health))
.route("/datasets", post(create_dataset))
.route("/datasets", get(list_datasets))
.route("/datasets/{id}", get(get_dataset))
.route("/datasets/by-name/{name}", get(get_dataset_by_name).delete(delete_dataset_by_name))
.route("/datasets/by-name/{name}/metadata", post(update_metadata))
.route("/datasets/by-name/{name}/resync", post(resync_dataset))
.route("/resync-missing", post(resync_all_missing))
.route("/migrate-buckets", post(migrate_buckets))
.route("/dedupe", post(dedupe_by_name))
// Phase D: AI-safe views
.route("/views", post(create_view).get(list_views))
.route("/views/{name}", get(get_view).delete(delete_view))
// Phase E: soft-delete tombstones
.route("/datasets/by-name/{name}/tombstone", post(tombstone_rows).get(list_tombstones))
// Phase 17: model profiles
.route("/profiles", post(create_profile).get(list_profiles))
.route("/profiles/{id}", get(get_profile).delete(delete_profile))
// Phase 41: profile-type routes
.route("/profiles/retrieval", get(list_retrieval_profiles))
.route("/profiles/memory", get(list_memory_profiles))
.route("/profiles/observer", get(list_observer_profiles))
.with_state(registry)
}
async fn health() -> &'static str {
"catalogd ok"
}
#[derive(Deserialize)]
struct CreateDatasetRequest {
name: String,
schema_fingerprint: String,
objects: Vec<ObjectRefRequest>,
}
#[derive(Deserialize)]
struct ObjectRefRequest {
bucket: String,
key: String,
size_bytes: u64,
}
#[derive(Serialize)]
struct DatasetResponse {
id: String,
name: String,
schema_fingerprint: String,
objects: Vec<ObjectRefResponse>,
created_at: String,
updated_at: String,
// Rich metadata
description: String,
owner: String,
sensitivity: Option<shared::types::Sensitivity>,
columns: Vec<shared::types::ColumnMeta>,
lineage: Option<shared::types::Lineage>,
freshness: Option<shared::types::FreshnessContract>,
tags: Vec<String>,
row_count: Option<u64>,
}
#[derive(Serialize)]
struct ObjectRefResponse {
bucket: String,
key: String,
size_bytes: u64,
created_at: String,
}
impl From<&shared::types::DatasetManifest> for DatasetResponse {
fn from(m: &shared::types::DatasetManifest) -> Self {
Self {
id: m.id.to_string(),
name: m.name.clone(),
schema_fingerprint: m.schema_fingerprint.0.clone(),
objects: m.objects.iter().map(|o| ObjectRefResponse {
bucket: o.bucket.clone(),
key: o.key.clone(),
size_bytes: o.size_bytes,
created_at: o.created_at.to_rfc3339(),
}).collect(),
created_at: m.created_at.to_rfc3339(),
updated_at: m.updated_at.to_rfc3339(),
description: m.description.clone(),
owner: m.owner.clone(),
sensitivity: m.sensitivity.clone(),
columns: m.columns.clone(),
lineage: m.lineage.clone(),
freshness: m.freshness.clone(),
tags: m.tags.clone(),
row_count: m.row_count,
}
}
}
async fn create_dataset(
State(registry): State<Registry>,
Json(req): Json<CreateDatasetRequest>,
) -> impl IntoResponse {
let now = chrono::Utc::now();
let objects: Vec<ObjectRef> = req.objects.into_iter().map(|o| ObjectRef {
bucket: o.bucket,
key: o.key,
size_bytes: o.size_bytes,
created_at: now,
}).collect();
match registry.register(req.name, SchemaFingerprint(req.schema_fingerprint), objects).await {
Ok(manifest) => {
let resp = DatasetResponse::from(&manifest);
Ok((StatusCode::CREATED, Json(resp)))
}
Err(e) => {
let status = if e.contains("different schema") {
StatusCode::CONFLICT
} else {
StatusCode::INTERNAL_SERVER_ERROR
};
Err((status, e))
}
}
}
async fn list_datasets(State(registry): State<Registry>) -> impl IntoResponse {
let datasets = registry.list().await;
let resp: Vec<DatasetResponse> = datasets.iter().map(DatasetResponse::from).collect();
Json(resp)
}
async fn get_dataset(
State(registry): State<Registry>,
Path(id): Path<String>,
) -> impl IntoResponse {
let uuid = Uuid::parse_str(&id).map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?;
let dataset_id = DatasetId(uuid);
match registry.get(&dataset_id).await {
Some(manifest) => Ok(Json(DatasetResponse::from(&manifest))),
None => Err((StatusCode::NOT_FOUND, format!("dataset not found: {id}"))),
}
}
async fn get_dataset_by_name(
State(registry): State<Registry>,
Path(name): Path<String>,
) -> impl IntoResponse {
match registry.get_by_name(&name).await {
Some(manifest) => Ok(Json(DatasetResponse::from(&manifest))),
None => Err((StatusCode::NOT_FOUND, format!("dataset not found: {name}"))),
}
}
/// Remove a dataset manifest by name. Metadata only — parquet files,
/// vector indexes, and tombstones are NOT cascade-deleted. See
/// `Registry::delete_dataset` for the full scope.
async fn delete_dataset_by_name(
State(registry): State<Registry>,
Path(name): Path<String>,
) -> impl IntoResponse {
match registry.delete_dataset(&name).await {
Ok(removed) => {
let body = serde_json::json!({ "name": name, "manifests_removed": removed });
Ok(Json(body))
}
Err(e) if e.starts_with("dataset not found") => {
Err((StatusCode::NOT_FOUND, e))
}
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
async fn update_metadata(
State(registry): State<Registry>,
Path(name): Path<String>,
Json(updates): Json<crate::registry::MetadataUpdate>,
) -> impl IntoResponse {
match registry.update_metadata(&name, updates).await {
Ok(manifest) => Ok(Json(DatasetResponse::from(&manifest))),
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
/// Re-read parquet footers for a single dataset and repopulate row_count
/// and columns from reality. Useful for repairing manifests whose metadata
/// was lost or never backfilled.
async fn resync_dataset(
State(registry): State<Registry>,
Path(name): Path<String>,
) -> impl IntoResponse {
match registry.resync_from_parquet(&name).await {
Ok(manifest) => Ok(Json(DatasetResponse::from(&manifest))),
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
#[derive(Serialize)]
struct ResyncAllResponse {
succeeded: Vec<ResyncOk>,
failed: Vec<ResyncErr>,
}
#[derive(Serialize)]
struct ResyncOk {
name: String,
row_count: u64,
}
#[derive(Serialize)]
struct ResyncErr {
name: String,
error: String,
}
/// Resync every dataset that currently has null row_count or empty columns.
async fn resync_all_missing(State(registry): State<Registry>) -> impl IntoResponse {
let (ok, err) = registry.resync_missing().await;
Json(ResyncAllResponse {
succeeded: ok.into_iter().map(|(name, row_count)| ResyncOk { name, row_count }).collect(),
failed: err.into_iter().map(|(name, error)| ResyncErr { name, error }).collect(),
})
}
/// Federation layer 2 one-shot: normalize every ObjectRef.bucket field
/// to the canonical "primary" value. Idempotent — re-running once
/// everything is canonical is a safe no-op.
async fn migrate_buckets(State(registry): State<Registry>) -> impl IntoResponse {
match registry.migrate_buckets_to_primary().await {
Ok(report) => Ok(Json(report)),
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
/// Collapse duplicate manifests by `name` — see `Registry::dedupe_by_name`.
/// Safe to run repeatedly; single-manifest datasets are untouched.
async fn dedupe_by_name(State(registry): State<Registry>) -> impl IntoResponse {
Json(registry.dedupe_by_name().await)
}
// --- Phase D: AI-safe views ---
#[derive(Deserialize)]
struct CreateViewRequest {
name: String,
base_dataset: String,
columns: Vec<String>,
#[serde(default)]
row_filter: Option<String>,
#[serde(default)]
column_redactions: std::collections::HashMap<String, shared::types::Redaction>,
#[serde(default)]
description: String,
#[serde(default)]
created_by: String,
}
async fn create_view(
State(registry): State<Registry>,
Json(req): Json<CreateViewRequest>,
) -> impl IntoResponse {
let view = shared::types::AiView {
name: req.name,
base_dataset: req.base_dataset,
columns: req.columns,
row_filter: req.row_filter,
column_redactions: req.column_redactions,
created_at: chrono::Utc::now(),
created_by: req.created_by,
description: req.description,
};
match registry.put_view(view).await {
Ok(v) => Ok((StatusCode::CREATED, Json(v))),
Err(e) => Err((StatusCode::BAD_REQUEST, e)),
}
}
async fn list_views(State(registry): State<Registry>) -> impl IntoResponse {
Json(registry.list_views().await)
}
async fn get_view(
State(registry): State<Registry>,
Path(name): Path<String>,
) -> impl IntoResponse {
match registry.get_view(&name).await {
Some(v) => Ok(Json(v)),
None => Err((StatusCode::NOT_FOUND, format!("view not found: {name}"))),
}
}
async fn delete_view(
State(registry): State<Registry>,
Path(name): Path<String>,
) -> impl IntoResponse {
match registry.delete_view(&name).await {
Ok(()) => Ok(StatusCode::NO_CONTENT),
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
// --- Phase E: soft-delete tombstones ---
#[derive(Deserialize)]
struct TombstoneRequest {
row_key_column: String,
row_key_values: Vec<String>,
#[serde(default)]
actor: String,
#[serde(default)]
reason: String,
}
#[derive(Serialize)]
struct TombstoneResponse {
dataset: String,
row_key_column: String,
rows_tombstoned: usize,
failures: Vec<String>,
}
async fn tombstone_rows(
State(registry): State<Registry>,
Path(name): Path<String>,
Json(req): Json<TombstoneRequest>,
) -> impl IntoResponse {
if req.row_key_values.is_empty() {
return Err((StatusCode::BAD_REQUEST, "row_key_values is empty".to_string()));
}
let mut ok = 0;
let mut failures = Vec::new();
for value in &req.row_key_values {
match registry
.add_tombstone(&name, &req.row_key_column, value, &req.actor, &req.reason)
.await
{
Ok(_) => ok += 1,
Err(e) => failures.push(format!("{value}: {e}")),
}
}
let status = if ok > 0 && failures.is_empty() {
StatusCode::CREATED
} else if ok > 0 {
StatusCode::MULTI_STATUS
} else {
StatusCode::BAD_REQUEST
};
Ok((status, Json(TombstoneResponse {
dataset: name,
row_key_column: req.row_key_column,
rows_tombstoned: ok,
failures,
})))
}
async fn list_tombstones(
State(registry): State<Registry>,
Path(name): Path<String>,
) -> impl IntoResponse {
match registry.list_tombstones(&name).await {
Ok(ts) => Ok(Json(ts)),
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
// --- Phase 17: Model profiles ---
#[derive(Deserialize)]
struct CreateProfileRequest {
id: String,
ollama_name: String,
#[serde(default)]
description: String,
bound_datasets: Vec<String>,
#[serde(default)]
hnsw_config: shared::types::ProfileHnswConfig,
#[serde(default = "default_embed_model_req")]
embed_model: String,
#[serde(default)]
created_by: String,
/// Federation: optional per-profile bucket (`profile:{id}` by convention).
/// Omitting keeps artifacts in primary.
#[serde(default)]
bucket: Option<String>,
/// ADR-019 hybrid: which vector backend to route this profile's
/// indexes to. Defaults to Parquet+HNSW.
#[serde(default)]
vector_backend: shared::types::VectorBackend,
/// Phase 41: Profile type for routing + activation behavior.
#[serde(default)]
profile_type: shared::types::ProfileType,
}
fn default_embed_model_req() -> String { "nomic-embed-text".to_string() }
async fn create_profile(
State(registry): State<Registry>,
Json(req): Json<CreateProfileRequest>,
) -> impl IntoResponse {
let profile = shared::types::ModelProfile {
id: req.id,
ollama_name: req.ollama_name,
description: req.description,
bound_datasets: req.bound_datasets,
hnsw_config: req.hnsw_config,
embed_model: req.embed_model,
created_at: chrono::Utc::now(),
created_by: req.created_by,
bucket: req.bucket,
vector_backend: req.vector_backend,
profile_type: req.profile_type,
};
match registry.put_profile(profile).await {
Ok(p) => Ok((StatusCode::CREATED, Json(p))),
Err(e) => Err((StatusCode::BAD_REQUEST, e)),
}
}
async fn list_profiles(State(registry): State<Registry>) -> impl IntoResponse {
Json(registry.list_profiles().await)
}
async fn list_retrieval_profiles(State(registry): State<Registry>) -> impl IntoResponse {
let all = registry.list_profiles().await;
let retrieval: Vec<_> = all.into_iter().filter(|p| p.profile_type == shared::types::ProfileType::Retrieval).collect();
Json(retrieval)
}
async fn list_memory_profiles(State(registry): State<Registry>) -> impl IntoResponse {
let all = registry.list_profiles().await;
let memory: Vec<_> = all.into_iter().filter(|p| p.profile_type == shared::types::ProfileType::Memory).collect();
Json(memory)
}
async fn list_observer_profiles(State(registry): State<Registry>) -> impl IntoResponse {
let all = registry.list_profiles().await;
let observer: Vec<_> = all.into_iter().filter(|p| p.profile_type == shared::types::ProfileType::Observer).collect();
Json(observer)
}
async fn get_profile(
State(registry): State<Registry>,
Path(id): Path<String>,
) -> impl IntoResponse {
match registry.get_profile(&id).await {
Some(p) => Ok(Json(p)),
None => Err((StatusCode::NOT_FOUND, format!("profile not found: {id}"))),
}
}
async fn delete_profile(
State(registry): State<Registry>,
Path(id): Path<String>,
) -> impl IntoResponse {
match registry.delete_profile(&id).await {
Ok(()) => Ok(StatusCode::NO_CONTENT),
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
+151
View File
@@ -0,0 +1,151 @@
//! Soft-delete tombstone storage (Phase E).
//!
//! One append-log per dataset at `_catalog/tombstones/{dataset}/batch_*.jsonl`.
//! Uses the shared `storaged::append_log::AppendLog` pattern so appends
//! are write-once (never rewrites existing files) and can be compacted.
//!
//! The store exposes a per-dataset cache of active tombstones for the
//! hot path (queryd filter construction) — that's why events are pulled
//! into an in-memory map on every call rather than scanning object
//! storage repeatedly.
use object_store::ObjectStore;
use shared::types::Tombstone;
use std::collections::HashMap;
use std::sync::Arc;
use storaged::append_log::{AppendLog, CompactStats};
use tokio::sync::RwLock;
const TOMBSTONE_PREFIX: &str = "_catalog/tombstones";
#[derive(Clone)]
pub struct TombstoneStore {
store: Arc<dyn ObjectStore>,
logs: Arc<RwLock<HashMap<String, Arc<AppendLog>>>>,
}
impl TombstoneStore {
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
Self {
store,
logs: Arc::new(RwLock::new(HashMap::new())),
}
}
fn prefix_for(dataset: &str) -> String {
// Sanitize dataset name for filesystem safety.
let safe: String = dataset
.chars()
.map(|c| if c.is_ascii_alphanumeric() || c == '_' || c == '-' { c } else { '_' })
.collect();
format!("{TOMBSTONE_PREFIX}/{}", safe)
}
async fn log_for(&self, dataset: &str) -> Arc<AppendLog> {
if let Some(log) = self.logs.read().await.get(dataset) {
return log.clone();
}
let mut guard = self.logs.write().await;
if let Some(log) = guard.get(dataset) {
return log.clone();
}
// Threshold of 1 — every tombstone is high-value. Compliance/audit
// doesn't tolerate "lost on restart"; we trade a small file count
// for guaranteed durability. Compaction merges later if volume grows.
let log = Arc::new(
AppendLog::new(self.store.clone(), Self::prefix_for(dataset))
.with_flush_threshold(1),
);
guard.insert(dataset.to_string(), log.clone());
log
}
/// Append one tombstone. Validates that the `row_key_column` matches
/// the column already used for this dataset (all tombstones for a
/// dataset share one key column so the query filter is well-defined).
/// Forces a flush so the tombstone is durable before this call returns.
pub async fn append(&self, ts: &Tombstone) -> Result<(), String> {
let existing = self.list(&ts.dataset).await?;
if let Some(prior) = existing.first() {
if prior.row_key_column != ts.row_key_column {
return Err(format!(
"dataset '{}' already uses '{}' as tombstone key; cannot mix with '{}'",
ts.dataset, prior.row_key_column, ts.row_key_column,
));
}
}
let line = serde_json::to_vec(ts).map_err(|e| e.to_string())?;
let log = self.log_for(&ts.dataset).await;
log.append(line).await?;
// Belt-and-suspenders: explicit flush in case the threshold is
// ever raised. Tombstones must be durable on return.
log.flush().await
}
/// All tombstones for a dataset (chronological).
pub async fn list(&self, dataset: &str) -> Result<Vec<Tombstone>, String> {
let log = self.log_for(dataset).await;
let lines = log.read_all().await?;
let mut out = Vec::with_capacity(lines.len());
for line in lines {
match serde_json::from_slice::<Tombstone>(&line) {
Ok(t) => out.push(t),
Err(e) => tracing::warn!("tombstones/{}: skip malformed entry: {e}", dataset),
}
}
Ok(out)
}
/// Per-dataset grouped view used by queryd — returns a map of
/// `{dataset -> (row_key_column, set_of_values)}` for every dataset
/// that has any tombstones.
pub async fn all_grouped(
&self,
datasets: &[String],
) -> Result<HashMap<String, (String, Vec<String>)>, String> {
let mut grouped = HashMap::new();
for dataset in datasets {
let ts = match self.list(dataset).await {
Ok(ts) => ts,
Err(_) => continue,
};
if ts.is_empty() { continue; }
let col = ts[0].row_key_column.clone();
let values: Vec<String> = ts.iter().map(|t| t.row_key_value.clone()).collect();
grouped.insert(dataset.clone(), (col, values));
}
Ok(grouped)
}
pub async fn compact(&self, dataset: &str) -> Result<CompactStats, String> {
let log = self.log_for(dataset).await;
log.compact().await
}
/// Remove every tombstone for a dataset. Called after a successful
/// parquet compaction has physically deleted those rows — the
/// tombstones have done their job and the journal can be cleared.
///
/// Implementation: drop the per-dataset AppendLog from the cache and
/// delete all its batch files. Next write starts fresh.
pub async fn clear(&self, dataset: &str) -> Result<usize, String> {
let prefix = format!("{}/", Self::prefix_for(dataset));
let keys = storaged::ops::list(&self.store, Some(&prefix)).await?;
let matching: Vec<String> = keys
.into_iter()
.filter(|k| {
let basename = k.rsplit('/').next().unwrap_or(k);
basename.starts_with("batch_") && basename.ends_with(".jsonl")
})
.collect();
let count = matching.len();
for key in &matching {
let _ = storaged::ops::delete(&self.store, key).await;
}
self.logs.write().await.remove(dataset);
if count > 0 {
tracing::info!("cleared {count} tombstone batch files for '{}'", dataset);
}
Ok(count)
}
}
+32
View File
@@ -0,0 +1,32 @@
[package]
name = "gateway"
version = "0.1.0"
edition = "2024"
[dependencies]
shared = { path = "../shared" }
storaged = { path = "../storaged" }
catalogd = { path = "../catalogd" }
queryd = { path = "../queryd" }
aibridge = { path = "../aibridge" }
ingestd = { path = "../ingestd" }
vectord = { path = "../vectord" }
journald = { path = "../journald" }
truth = { path = "../truth" }
tokio = { workspace = true }
axum = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }
tracing = { workspace = true }
tracing-subscriber = { workspace = true }
tower-http = { workspace = true }
object_store = { workspace = true }
proto = { path = "../proto" }
tonic = { workspace = true }
opentelemetry = { workspace = true }
opentelemetry_sdk = { workspace = true }
opentelemetry-stdout = { workspace = true }
tracing-opentelemetry = { workspace = true }
arrow = { workspace = true }
chrono = { workspace = true }
reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] }
+160
View File
@@ -0,0 +1,160 @@
/// Access control: field-level sensitivity enforcement, column masking, query audit.
/// Evaluates policies at query time — not at storage level.
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use shared::types::Sensitivity;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
/// An agent's role determines what sensitivity levels they can see.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentRole {
pub agent_name: String,
pub role: String, // "recruiter", "admin", "analyst", "agent"
pub allowed_sensitivity: Vec<Sensitivity>, // what they can see unmasked
}
/// A query audit entry.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryAudit {
pub id: String,
pub agent: String,
pub sql: String,
pub datasets_accessed: Vec<String>,
pub pii_fields_accessed: Vec<String>,
pub timestamp: DateTime<Utc>,
pub row_count: usize,
pub allowed: bool,
pub masked_fields: Vec<String>,
}
/// Access control manager.
#[derive(Clone)]
pub struct AccessControl {
roles: Arc<RwLock<HashMap<String, AgentRole>>>,
audit_log: Arc<RwLock<Vec<QueryAudit>>>,
enabled: bool,
}
impl AccessControl {
pub fn new(enabled: bool) -> Self {
let ac = Self {
roles: Arc::new(RwLock::new(HashMap::new())),
audit_log: Arc::new(RwLock::new(Vec::new())),
enabled,
};
ac
}
/// Register default roles.
pub async fn register_defaults(&self) {
let defaults = vec![
AgentRole {
agent_name: "admin".into(),
role: "admin".into(),
allowed_sensitivity: vec![
Sensitivity::Public, Sensitivity::Internal,
Sensitivity::Pii, Sensitivity::Phi, Sensitivity::Financial,
],
},
AgentRole {
agent_name: "recruiter".into(),
role: "recruiter".into(),
allowed_sensitivity: vec![
Sensitivity::Public, Sensitivity::Internal, Sensitivity::Pii,
],
},
AgentRole {
agent_name: "analyst".into(),
role: "analyst".into(),
allowed_sensitivity: vec![
Sensitivity::Public, Sensitivity::Internal, Sensitivity::Financial,
],
},
AgentRole {
agent_name: "agent".into(),
role: "agent".into(),
allowed_sensitivity: vec![
Sensitivity::Public, Sensitivity::Internal,
],
},
];
let mut roles = self.roles.write().await;
for role in defaults {
roles.insert(role.agent_name.clone(), role);
}
}
/// Register or update an agent role.
pub async fn set_role(&self, role: AgentRole) {
self.roles.write().await.insert(role.agent_name.clone(), role);
}
/// Get an agent's role. Called by `GET /access/roles/{agent}`.
pub async fn get_role(&self, agent: &str) -> Option<AgentRole> {
self.roles.read().await.get(agent).cloned()
}
/// List all roles.
pub async fn list_roles(&self) -> Vec<AgentRole> {
self.roles.read().await.values().cloned().collect()
}
/// Check if an agent can see a field with given sensitivity.
pub async fn can_access(&self, agent: &str, sensitivity: &Sensitivity) -> bool {
if !self.enabled { return true; }
match self.roles.read().await.get(agent) {
Some(role) => role.allowed_sensitivity.contains(sensitivity),
None => false, // unknown agent = no access
}
}
/// Determine which fields should be masked for an agent.
#[allow(dead_code)]
pub async fn masked_fields(
&self,
agent: &str,
columns: &[shared::types::ColumnMeta],
) -> Vec<String> {
if !self.enabled { return vec![]; }
let role = match self.roles.read().await.get(agent) {
Some(r) => r.clone(),
None => return columns.iter().filter(|c| c.sensitivity.is_some()).map(|c| c.name.clone()).collect(),
};
columns.iter()
.filter(|col| {
if let Some(ref sens) = col.sensitivity {
!role.allowed_sensitivity.contains(sens)
} else {
false
}
})
.map(|col| col.name.clone())
.collect()
}
/// Log a query for audit.
#[allow(dead_code)]
pub async fn log_query(&self, audit: QueryAudit) {
self.audit_log.write().await.push(audit);
}
/// Get recent audit entries.
pub async fn recent_audit(&self, limit: usize) -> Vec<QueryAudit> {
let log = self.audit_log.read().await;
let start = log.len().saturating_sub(limit);
log[start..].iter().rev().cloned().collect()
}
/// Reports whether access-control enforcement is active.
/// Called by `GET /access/enabled` — ops tooling / dashboards poll
/// this to confirm the auth posture of the running gateway.
pub fn is_enabled(&self) -> bool {
self.enabled
}
}
+82
View File
@@ -0,0 +1,82 @@
use axum::{
Json, Router,
extract::{Path, Query, State},
http::StatusCode,
response::IntoResponse,
routing::{get, post},
};
use serde::Deserialize;
use crate::access::{AccessControl, AgentRole};
pub fn router(ac: AccessControl) -> Router {
Router::new()
.route("/roles", get(list_roles))
.route("/roles", post(set_role))
// Scrum iter 11 / P13-001 finding: get_role was #[allow(dead_code)]
// because nothing called it — dead until exposed. Route activates it.
// Returns 404 when the agent isn't registered so clients can
// distinguish "missing role" from "access denied."
.route("/roles/{agent}", get(get_role))
.route("/enabled", get(enabled_status))
.route("/audit", get(query_audit))
.route("/check", post(check_access))
.with_state(ac)
}
async fn list_roles(State(ac): State<AccessControl>) -> impl IntoResponse {
Json(ac.list_roles().await)
}
async fn set_role(
State(ac): State<AccessControl>,
Json(role): Json<AgentRole>,
) -> impl IntoResponse {
let name = role.agent_name.clone();
ac.set_role(role).await;
(StatusCode::OK, format!("role set: {name}"))
}
#[derive(Deserialize)]
struct AuditQuery {
limit: Option<usize>,
}
async fn query_audit(
State(ac): State<AccessControl>,
Query(q): Query<AuditQuery>,
) -> impl IntoResponse {
Json(ac.recent_audit(q.limit.unwrap_or(50)).await)
}
#[derive(Deserialize)]
struct CheckRequest {
agent: String,
sensitivity: shared::types::Sensitivity,
}
async fn check_access(
State(ac): State<AccessControl>,
Json(req): Json<CheckRequest>,
) -> impl IntoResponse {
let allowed = ac.can_access(&req.agent, &req.sensitivity).await;
Json(serde_json::json!({
"agent": req.agent,
"sensitivity": req.sensitivity,
"allowed": allowed,
}))
}
async fn get_role(
State(ac): State<AccessControl>,
Path(agent): Path<String>,
) -> impl IntoResponse {
match ac.get_role(&agent).await {
Some(role) => Ok(Json(role)),
None => Err((StatusCode::NOT_FOUND, format!("no role registered for agent '{agent}'"))),
}
}
async fn enabled_status(State(ac): State<AccessControl>) -> impl IntoResponse {
Json(serde_json::json!({ "enabled": ac.is_enabled() }))
}
+57
View File
@@ -0,0 +1,57 @@
use axum::{
extract::Request,
http::StatusCode,
middleware::Next,
response::Response,
};
// API key auth middleware. Checks X-API-Key header against configured key.
// Fixed P5-001 (2026-04-23): previously #[allow(dead_code)] — the function
// existed but was never layered onto the router, so [auth] enabled=true
// silently enforced nothing. Now wired via from_fn_with_state in main.rs.
pub async fn api_key_auth(
axum::extract::State(expected): axum::extract::State<ApiKey>,
request: Request,
next: Next,
) -> Result<Response, StatusCode> {
// /health stays public (LB/systemd probes). Every other route is gated.
if request.uri().path() == "/health" {
return Ok(next.run(request).await);
}
let provided = request
.headers()
.get("x-api-key")
.and_then(|v| v.to_str().ok());
// Constant-time-ish eq on the raw bytes; good enough for a shared-secret
// X-API-Key. Timing-attack resistance here matters less than the
// equivalent HMAC check would; adopt subtle crate if key-space grows.
match provided {
Some(key) if eq_ct(key.as_bytes(), expected.0.as_bytes()) => {
Ok(next.run(request).await)
}
_ => {
tracing::warn!(
path = %request.uri().path(),
"unauthorized request: missing or invalid API key",
);
Err(StatusCode::UNAUTHORIZED)
}
}
}
fn eq_ct(a: &[u8], b: &[u8]) -> bool {
if a.len() != b.len() {
return false;
}
let mut diff: u8 = 0;
for (x, y) in a.iter().zip(b.iter()) {
diff |= x ^ y;
}
diff == 0
}
/// Wrapper type for the API key, stored in request extensions.
#[derive(Clone)]
pub struct ApiKey(pub String);
@@ -0,0 +1,388 @@
//! KB context loader — reads recent signal from `data/_kb/*.jsonl` for
//! a given sig_hash + task_class and returns a compact summary.
//!
//! This is the "pipe to the overviewer" from the 2026-04-23 session:
//! the overseer tier (T3, gpt-oss:120b) consumes this context before
//! generating a correction, so its suggestions are informed by
//! historical cost / latency / outcome / prior-correction patterns
//! across ALL profiles that have run this task class — not just the
//! single current loop.
//!
//! Hot-swap profiles read the SAME pool. When a profile activates and
//! starts iterating, its KB context is the shared surface — one
//! profile's learning becomes every profile's starting point.
//!
//! Best-effort throughout: missing files, corrupt rows, empty
//! directories all produce an empty KbContext. The overseer works
//! fine with no history; we just can't seed it then.
use serde::Serialize;
use std::path::Path;
use tokio::io::AsyncBufReadExt;
/// Compact summary returned to the overseer. Bounded size — recent
/// outcomes + corrections plus rolled-up rates. Goal is to fit in a
/// prompt without eating the overseer's context budget.
#[derive(Debug, Clone, Default, Serialize)]
pub struct KbContext {
pub sig_hash: String,
pub task_class: String,
pub recent_outcomes: Vec<OutcomeSummary>,
pub recent_corrections: Vec<CorrectionSummary>,
pub success_rate: Option<f64>,
pub avg_turns: Option<f64>,
pub avg_latency_ms: Option<u64>,
pub total_observed: u32,
}
#[derive(Debug, Clone, Serialize)]
pub struct OutcomeSummary {
pub created_at: String,
pub ok: bool,
pub polarity: String,
pub turns: u32,
pub latency_ms: u64,
pub total_tokens: u64,
pub error: Option<String>,
}
#[derive(Debug, Clone, Serialize)]
pub struct CorrectionSummary {
pub created_at: String,
pub reason: String,
pub correction_preview: String, // first 300 chars
pub applied_at_turn: u32,
}
const OUTCOMES_PATH: &str = "data/_kb/outcomes.jsonl";
const CORRECTIONS_PATH: &str = "data/_kb/overseer_corrections.jsonl";
const RECENT_OUTCOME_LIMIT: usize = 5;
const RECENT_CORRECTION_LIMIT: usize = 3;
const AGGREGATE_WINDOW: usize = 50;
impl KbContext {
/// Build context from the default KB paths.
pub async fn load_for(sig_hash: &str, task_class: &str) -> Self {
Self::load_from(
sig_hash, task_class,
Path::new(OUTCOMES_PATH), Path::new(CORRECTIONS_PATH),
).await
}
/// Path-taking variant — tests inject tmp files without touching
/// the real KB directory (same pattern as append_outcomes_row_at).
pub async fn load_from(
sig_hash: &str,
task_class: &str,
outcomes_path: &Path,
corrections_path: &Path,
) -> Self {
let mut ctx = KbContext {
sig_hash: sig_hash.to_string(),
task_class: task_class.to_string(),
..Default::default()
};
// Scan outcomes — matches on sig_hash primary, task_class
// secondary (so different geos for the same task_class still
// contribute to aggregate rates even though they won't make
// the top-5 recent). The bounded window keeps scan cost
// linear in file size — we're reading tail only.
let outcome_rows = tail_matching(
outcomes_path, AGGREGATE_WINDOW * 4,
|row| {
let row_sig = row.get("sig_hash").and_then(|v| v.as_str()).unwrap_or("");
let row_tc = row.get("task_class").and_then(|v| v.as_str()).unwrap_or("");
row_sig == sig_hash || row_tc == task_class
},
).await;
// Recent outcomes: exact sig_hash match first (strongest
// signal), then task_class fallback up to the limit.
let mut exact: Vec<OutcomeSummary> = Vec::new();
let mut loose: Vec<OutcomeSummary> = Vec::new();
for row in &outcome_rows {
let row_sig = row.get("sig_hash").and_then(|v| v.as_str()).unwrap_or("");
let summary = summarize_outcome(row);
if row_sig == sig_hash { exact.push(summary); }
else { loose.push(summary); }
}
ctx.recent_outcomes = exact.into_iter().rev().take(RECENT_OUTCOME_LIMIT).collect();
if ctx.recent_outcomes.len() < RECENT_OUTCOME_LIMIT {
let need = RECENT_OUTCOME_LIMIT - ctx.recent_outcomes.len();
ctx.recent_outcomes.extend(loose.into_iter().rev().take(need));
}
// Aggregate rates across the full matched window (both
// sig_hash and task_class matches — gives a stable rate even
// on sparse sig_hash history).
let window = outcome_rows.iter().rev().take(AGGREGATE_WINDOW);
let mut ok_count = 0u32;
let mut total = 0u32;
let mut turn_sum = 0u32;
let mut latency_sum = 0u64;
for row in window {
total += 1;
if row.get("ok").and_then(|v| v.as_bool()).unwrap_or(false) { ok_count += 1; }
turn_sum += row.get("turns").and_then(|v| v.as_u64()).unwrap_or(0) as u32;
latency_sum += row.get("usage")
.and_then(|u| u.get("latency_ms"))
.and_then(|v| v.as_u64()).unwrap_or(0);
}
if total > 0 {
ctx.total_observed = total;
ctx.success_rate = Some(ok_count as f64 / total as f64);
ctx.avg_turns = Some(turn_sum as f64 / total as f64);
ctx.avg_latency_ms = Some(latency_sum / total as u64);
}
// Overseer corrections. Prefer sig_hash match; fall back to
// task_class. The overseer reading its OWN prior corrections
// is the main point — if the last 3 attempts produced
// corrections X, Y, Z, the new correction should acknowledge
// those patterns rather than suggest X for the fourth time.
let correction_rows = tail_matching(
corrections_path, RECENT_CORRECTION_LIMIT * 4,
|row| {
let row_sig = row.get("sig_hash").and_then(|v| v.as_str()).unwrap_or("");
let row_tc = row.get("task_class").and_then(|v| v.as_str()).unwrap_or("");
row_sig == sig_hash || row_tc == task_class
},
).await;
let mut c_exact: Vec<CorrectionSummary> = Vec::new();
let mut c_loose: Vec<CorrectionSummary> = Vec::new();
for row in &correction_rows {
let row_sig = row.get("sig_hash").and_then(|v| v.as_str()).unwrap_or("");
let summary = summarize_correction(row);
if row_sig == sig_hash { c_exact.push(summary); }
else { c_loose.push(summary); }
}
ctx.recent_corrections = c_exact.into_iter().rev().take(RECENT_CORRECTION_LIMIT).collect();
if ctx.recent_corrections.len() < RECENT_CORRECTION_LIMIT {
let need = RECENT_CORRECTION_LIMIT - ctx.recent_corrections.len();
ctx.recent_corrections.extend(c_loose.into_iter().rev().take(need));
}
ctx
}
/// Compact string form for the overseer prompt. Deterministic
/// ordering + bounded length so prompt caching stays stable
/// across iterations on the same task.
pub fn to_prompt_section(&self) -> String {
let mut s = String::new();
s.push_str("## Knowledge Base Context\n");
if let (Some(rate), Some(turns), Some(lat)) = (self.success_rate, self.avg_turns, self.avg_latency_ms) {
s.push_str(&format!(
"Across {} prior similar runs: success_rate={:.1}%, avg_turns={:.1}, avg_latency_ms={}\n",
self.total_observed, rate * 100.0, turns, lat,
));
} else {
s.push_str("No prior similar runs recorded.\n");
}
if !self.recent_outcomes.is_empty() {
s.push_str(&format!("\nRecent {} outcomes:\n", self.recent_outcomes.len()));
for o in &self.recent_outcomes {
let err = o.error.as_deref().map(|e| format!(" — {}", truncate(e, 80))).unwrap_or_default();
s.push_str(&format!(
" [{}] ok={} turns={} tokens={} lat={}ms{}\n",
&o.created_at[..19.min(o.created_at.len())],
o.ok, o.turns, o.total_tokens, o.latency_ms, err,
));
}
}
if !self.recent_corrections.is_empty() {
s.push_str(&format!("\nRecent {} overseer corrections (yours — don't repeat):\n", self.recent_corrections.len()));
for c in &self.recent_corrections {
s.push_str(&format!(
" [{}] turn={} reason={} correction={}\n",
&c.created_at[..19.min(c.created_at.len())],
c.applied_at_turn,
truncate(&c.reason, 40),
truncate(&c.correction_preview, 200),
));
}
}
s
}
}
fn summarize_outcome(row: &serde_json::Value) -> OutcomeSummary {
OutcomeSummary {
created_at: row.get("created_at").and_then(|v| v.as_str()).unwrap_or("").to_string(),
ok: row.get("ok").and_then(|v| v.as_bool()).unwrap_or(false),
polarity: row.get("polarity").and_then(|v| v.as_str()).unwrap_or("").to_string(),
turns: row.get("turns").and_then(|v| v.as_u64()).unwrap_or(0) as u32,
latency_ms: row.get("usage").and_then(|u| u.get("latency_ms"))
.and_then(|v| v.as_u64()).unwrap_or(0),
total_tokens: row.get("usage").and_then(|u| u.get("total_tokens"))
.and_then(|v| v.as_u64()).unwrap_or(0),
error: row.get("error").and_then(|v| v.as_str()).map(String::from),
}
}
fn summarize_correction(row: &serde_json::Value) -> CorrectionSummary {
let preview = row.get("correction").and_then(|v| v.as_str()).unwrap_or("");
CorrectionSummary {
created_at: row.get("created_at").and_then(|v| v.as_str()).unwrap_or("").to_string(),
reason: row.get("reason").and_then(|v| v.as_str()).unwrap_or("").to_string(),
correction_preview: truncate(preview, 300),
applied_at_turn: row.get("applied_at_turn").and_then(|v| v.as_u64()).unwrap_or(0) as u32,
}
}
fn truncate(s: &str, n: usize) -> String {
if s.len() <= n { s.to_string() } else { format!("{}…", &s[..n]) }
}
/// Read a JSONL file from the tail, returning at most `limit` rows
/// that match `filter`. Missing file returns empty. Corrupt lines are
/// skipped. Limit is honored from the tail — a full-file scan with an
/// in-memory ring would be wasteful for large outcomes histories, but
/// we cap at reading the whole file and filtering post-hoc for now
/// (reverse-seek line iteration is a real engineering task and the
/// file is bounded by ingest rate; revisit when it bites).
async fn tail_matching<F>(
path: &Path,
limit: usize,
filter: F,
) -> Vec<serde_json::Value>
where
F: Fn(&serde_json::Value) -> bool,
{
let Ok(file) = tokio::fs::File::open(path).await else { return Vec::new(); };
let reader = tokio::io::BufReader::new(file);
let mut lines = reader.lines();
let mut matches: Vec<serde_json::Value> = Vec::new();
while let Ok(Some(line)) = lines.next_line().await {
let Ok(v) = serde_json::from_str::<serde_json::Value>(&line) else { continue };
if filter(&v) {
matches.push(v);
if matches.len() > limit {
// Keep the most-recent window only — drop from the
// front as we go rather than buffering everything.
matches.remove(0);
}
}
}
matches
}
#[cfg(test)]
mod tests {
use super::*;
use tokio::io::AsyncWriteExt;
async fn write_fixture(path: &Path, rows: Vec<serde_json::Value>) {
if let Some(dir) = path.parent() {
tokio::fs::create_dir_all(dir).await.unwrap();
}
let mut f = tokio::fs::OpenOptions::new()
.create(true).write(true).truncate(true).open(path).await.unwrap();
for r in rows {
let mut line = serde_json::to_string(&r).unwrap();
line.push('\n');
f.write_all(line.as_bytes()).await.unwrap();
}
}
fn tmp_path(name: &str) -> std::path::PathBuf {
let nanos = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_nanos();
std::env::temp_dir().join(format!("lh_kb_ctx_{}_{}_{}", std::process::id(), nanos, name))
}
#[tokio::test]
async fn empty_files_produce_empty_context() {
let op = tmp_path("outcomes.jsonl");
let cp = tmp_path("corrections.jsonl");
let ctx = KbContext::load_from("sig123", "staffing.fill", &op, &cp).await;
assert!(ctx.recent_outcomes.is_empty());
assert!(ctx.recent_corrections.is_empty());
assert!(ctx.success_rate.is_none());
assert_eq!(ctx.total_observed, 0);
}
#[tokio::test]
async fn exact_sig_hash_matches_take_priority() {
let op = tmp_path("outcomes.jsonl");
let cp = tmp_path("corrections.jsonl");
write_fixture(&op, vec![
// Other sig_hash, same task_class — loose match
serde_json::json!({
"sig_hash": "other", "task_class": "staffing.fill",
"ok": false, "polarity": "failure_pattern", "turns": 1,
"usage": {"latency_ms": 1000, "total_tokens": 100},
"created_at": "2026-04-22T10:00:00Z",
}),
// Exact sig_hash — should lead
serde_json::json!({
"sig_hash": "sig123", "task_class": "staffing.fill",
"ok": true, "polarity": "success_confirmation", "turns": 3,
"usage": {"latency_ms": 2000, "total_tokens": 500},
"created_at": "2026-04-23T10:00:00Z",
}),
]).await;
write_fixture(&cp, vec![]).await;
let ctx = KbContext::load_from("sig123", "staffing.fill", &op, &cp).await;
assert_eq!(ctx.recent_outcomes.len(), 2);
assert_eq!(ctx.recent_outcomes[0].created_at, "2026-04-23T10:00:00Z");
assert_eq!(ctx.recent_outcomes[0].ok, true);
assert_eq!(ctx.total_observed, 2);
assert!((ctx.success_rate.unwrap() - 0.5).abs() < 0.001);
}
#[tokio::test]
async fn corrupt_rows_are_skipped() {
let op = tmp_path("outcomes.jsonl");
let cp = tmp_path("corrections.jsonl");
// Mix valid + invalid — invalid should be silently skipped.
if let Some(dir) = op.parent() { tokio::fs::create_dir_all(dir).await.unwrap(); }
tokio::fs::write(&op, "not json\n{\"sig_hash\":\"sig1\",\"task_class\":\"tc\",\"ok\":true,\"turns\":1,\"usage\":{}}\ngarbage\n").await.unwrap();
write_fixture(&cp, vec![]).await;
let ctx = KbContext::load_from("sig1", "tc", &op, &cp).await;
assert_eq!(ctx.recent_outcomes.len(), 1);
}
#[tokio::test]
async fn corrections_preview_is_truncated() {
let op = tmp_path("outcomes.jsonl");
let cp = tmp_path("corrections.jsonl");
let long = "x".repeat(500);
write_fixture(&op, vec![]).await;
write_fixture(&cp, vec![serde_json::json!({
"sig_hash": "sig1", "task_class": "tc",
"reason": "abort", "correction": long, "applied_at_turn": 3,
"created_at": "2026-04-23T10:00:00Z",
})]).await;
let ctx = KbContext::load_from("sig1", "tc", &op, &cp).await;
assert_eq!(ctx.recent_corrections.len(), 1);
// 300-char cap + 3-byte UTF-8 ellipsis character = 303-byte worst case.
assert!(ctx.recent_corrections[0].correction_preview.len() <= 303);
}
#[test]
fn prompt_section_is_stable_for_empty_context() {
let ctx = KbContext::default();
let s = ctx.to_prompt_section();
assert!(s.contains("No prior similar runs recorded"));
}
#[test]
fn prompt_section_reports_aggregate_rates() {
let ctx = KbContext {
total_observed: 10,
success_rate: Some(0.7),
avg_turns: Some(4.2),
avg_latency_ms: Some(45000),
..Default::default()
};
let s = ctx.to_prompt_section();
assert!(s.contains("success_rate=70.0%"));
assert!(s.contains("avg_turns=4.2"));
assert!(s.contains("avg_latency_ms=45000"));
}
}
File diff suppressed because it is too large Load Diff
+345
View File
@@ -0,0 +1,345 @@
mod access;
mod access_service;
mod auth;
mod execution_loop;
mod observability;
mod tools;
mod v1;
use axum::{Router, extract::DefaultBodyLimit, routing::get};
use proto::lakehouse::catalog_service_server::CatalogServiceServer;
use shared::config::Config;
use tower_http::cors::{Any, CorsLayer};
use tower_http::trace::TraceLayer;
#[tokio::main]
async fn main() {
// Load config
let config = Config::load_or_default();
// Initialize tracing + observability
observability::init_tracing(
&config.observability.service_name,
&config.observability.exporter,
);
tracing::info!("config loaded: gateway={}:{}, storage={}",
config.gateway.host, config.gateway.port, config.storage.root);
// Secrets provider
let secrets = std::sync::Arc::new(shared::secrets::FileSecretsProvider::from_env());
if let Err(e) = secrets.load().await {
tracing::warn!("secrets load: {e} — bucket credentials may be missing");
}
// Federation: bucket registry (primary + any [[storage.buckets]] entries + rescue)
let bucket_registry = std::sync::Arc::new(
storaged::registry::BucketRegistry::from_config(&config.storage, secrets.clone())
.await
.expect("bucket registry init failed")
);
let buckets = bucket_registry.list().await;
tracing::info!(
"bucket registry: {} buckets [{}], rescue={:?}",
buckets.len(),
buckets.iter().map(|b| format!("{}:{}", b.name, b.backend)).collect::<Vec<_>>().join(", "),
bucket_registry.rescue_name(),
);
// Back-compat: existing components that expect a single Arc<ObjectStore>
// get the primary bucket's store. As we migrate call sites to be bucket-
// aware, they'll call bucket_registry.get(name) directly.
let store = bucket_registry.default_store();
// Catalog
let registry = catalogd::registry::Registry::new(store.clone());
if let Err(e) = registry.rebuild().await {
tracing::warn!("catalog rebuild failed (empty store?): {e}");
}
// Query engine with 16GB memory cache
let cache = queryd::cache::MemCache::new(16 * 1024 * 1024 * 1024);
let engine = queryd::context::QueryEngine::new(registry.clone(), bucket_registry.clone(), cache);
// Event journal — append-only mutation log (flush every 100 events)
let journal = journald::journal::Journal::new(store.clone(), 100);
// Access control
let access = access::AccessControl::new(config.auth.enabled);
access.register_defaults().await;
// Phase 42 — file-backed truth rules. Probes the `truth/` directory
// at repo root (or $LAKEHOUSE_TRUTH_DIR override) and logs how many
// rules load. Current request paths still build their own stores
// via truth::default_truth_store() / truth::sql_query_guard_store();
// the composed-at-boot store gets plumbed through V1State in a
// follow-up. This boot probe catches parse errors + duplicate-ID
// collisions early rather than at first request.
{
let truth_dir = std::env::var("LAKEHOUSE_TRUTH_DIR")
.unwrap_or_else(|_| "/home/profit/lakehouse/truth".to_string());
if std::path::Path::new(&truth_dir).exists() {
let mut probe_store = truth::default_truth_store();
match truth::loader::load_from_dir(&mut probe_store, &truth_dir) {
Ok(n) => tracing::info!("truth: loaded {n} file-backed rule(s) from {truth_dir}"),
Err(e) => tracing::warn!("truth: failed to load rules from {truth_dir}: {e}"),
}
} else {
tracing::debug!("truth: no rule dir at {truth_dir}, skipping file-backed load");
}
}
// Workspace manager for agent-specific overlays
let workspace_mgr = queryd::workspace::WorkspaceManager::new(store.clone());
if let Err(e) = workspace_mgr.rebuild().await {
tracing::warn!("workspace rebuild: {e}");
}
// AI sidecar client
let ai_client = aibridge::client::AiClient::new(&config.sidecar.url);
// Vector service components — built before the router because both the
// /vectors service AND ingestd need the agent handle to enqueue triggers.
let index_reg = vectord::index_registry::IndexRegistry::new(store.clone());
let _ = index_reg.rebuild().await;
let hnsw = vectord::hnsw::HnswStore::new();
let emb_cache = vectord::embedding_cache::EmbeddingCache::new(store.clone());
// Phase B (federation layer 2): trial journals + promotion files now
// live in each index's recorded bucket (IndexMeta.bucket), not always
// primary. Journal / registry resolve per-call via the bucket registry.
let tj = vectord::trial::TrialJournal::new(bucket_registry.clone(), index_reg.clone());
let pr = vectord::promotion::PromotionRegistry::new(bucket_registry.clone(), index_reg.clone());
let hs = vectord::harness::HarnessStore::new(bucket_registry.clone(), index_reg.clone());
// Phase 19: playbook memory. Load cached state; empty-on-miss is fine,
// operators call POST /vectors/playbook_memory/rebuild to populate.
let pbm = vectord::playbook_memory::PlaybookMemory::new(store.clone());
let _ = pbm.load_from_storage().await;
// Pathway memory — consensus-designed sidecar for full-context
// backtracking + hot-swap of successful review pathways. Same
// load-on-boot pattern as playbook_memory: empty state is fine,
// operators start populating via scrum_master_pipeline.ts.
let pwm = vectord::pathway_memory::PathwayMemory::new(store.clone());
let _ = pwm.load_from_storage().await;
// Phase 16.2: spawn the autotune agent. When config.agent.enabled=false
// this returns a handle that drops triggers silently — no surprise load.
let agent_cfg = vectord::agent::AgentConfig {
enabled: config.agent.enabled,
cycle_interval_secs: config.agent.cycle_interval_secs,
cooldown_between_trials_secs: config.agent.cooldown_between_trials_secs,
min_recall: config.agent.min_recall,
max_trials_per_hour: config.agent.max_trials_per_hour,
};
let agent_handle = vectord::agent::spawn(
agent_cfg,
vectord::agent::AgentDeps {
store: store.clone(),
ai_client: ai_client.clone(),
catalog: registry.clone(),
index_registry: index_reg.clone(),
hnsw_store: hnsw.clone(),
embedding_cache: emb_cache.clone(),
trial_journal: tj.clone(),
promotion_registry: pr.clone(),
harness_store: hs.clone(),
},
);
// HTTP router
let mut app = Router::new()
.route("/health", get(health))
.nest("/storage", storaged::service::router(store.clone())
.merge(storaged::federation_service::router(bucket_registry.clone())))
.nest("/catalog", catalogd::service::router(registry.clone()))
.nest("/query", queryd::service::router(engine.clone()))
.nest("/ai", aibridge::service::router(ai_client.clone()))
;
// Phase E: scheduled ingest
let sched_store = ingestd::schedule::ScheduleStore::new(store.clone());
if let Ok(n) = sched_store.rebuild().await {
if n > 0 { tracing::info!("rebuilt {n} persisted schedule(s)"); }
}
let scheduler = ingestd::schedule::Scheduler {
store: sched_store.clone(),
ingest: ingestd::schedule::SchedulerIngestDeps {
store: store.clone(),
registry: registry.clone(),
buckets: bucket_registry.clone(),
agent_handle: Some(agent_handle.clone()),
index_registry: index_reg.clone(),
},
tick_secs: 10,
};
scheduler.spawn();
app = app
.nest("/ingest", ingestd::service::router(ingestd::service::IngestState {
store: store.clone(),
registry: registry.clone(),
buckets: bucket_registry.clone(),
agent_handle: agent_handle.clone(),
index_registry: index_reg.clone(),
schedules: sched_store,
// P9-001 fix 2026-04-23: journal reference flows into ingest so
// successful uploads emit a record_ingest event. Journal is Clone
// (Arc<RwLock> inside) so the /journal nest below still sees the
// same buffer + persistence.
journal: Some(journal.clone()),
}))
.nest("/vectors", vectord::service::router(vectord::service::VectorState {
store: store.clone(),
ai_client: ai_client.clone(),
job_tracker: vectord::jobs::JobTracker::new(),
index_registry: index_reg.clone(),
hnsw_store: hnsw,
embedding_cache: emb_cache,
trial_journal: tj,
harness_store: hs,
catalog: registry.clone(),
promotion_registry: pr,
agent_handle,
bucket_registry: bucket_registry.clone(),
active_profile: std::sync::Arc::new(tokio::sync::RwLock::new(None)),
lance: vectord::lance_backend::LanceRegistry::new(
bucket_registry.clone(), index_reg.clone(),
),
playbook_memory: pbm,
pathway_memory: pwm,
embed_semaphore: std::sync::Arc::new(tokio::sync::Semaphore::new(1)),
}))
.nest("/workspaces", queryd::workspace_service::router(workspace_mgr))
.nest("/journal", journald::service::router(journal))
.nest("/access", access_service::router(access))
.nest("/tools", tools::service::router({
let tool_reg = tools::registry::ToolRegistry::new();
tool_reg.register_defaults().await;
tools::ToolState {
registry: tool_reg,
query_fn: tools::QueryExecutor::new(engine.clone()),
truth: std::sync::Arc::new(truth::sql_query_guard_store()),
}
}))
// Phase 38 — Universal API skeleton. Thin OpenAI-compatible
// surface over the existing aibridge → Ollama path. Future
// phases add provider adapters (39), routing engine (40),
// session state (41), etc. All without changing this mount.
.nest("/v1", v1::router(v1::V1State {
ai_client: ai_client.clone(),
usage: std::sync::Arc::new(tokio::sync::RwLock::new(v1::Usage::default())),
// Phase 39 first slice — Ollama Cloud adapter. Key resolved
// from OLLAMA_CLOUD_KEY env, then /root/llm_team_config.json,
// then OLLAMA_CLOUD_API_KEY env. None = cloud routes 503.
ollama_cloud_key: {
let k = v1::ollama_cloud::resolve_cloud_key();
if k.is_some() {
tracing::info!("v1: Ollama Cloud key loaded — /v1/chat provider=ollama_cloud enabled");
} else {
tracing::warn!("v1: no Ollama Cloud key in env or /root/llm_team_config.json — cloud routes will 503");
}
k
},
openrouter_key: {
// 2026-04-24 free-tier rescue rung for iter 5+. Shares
// the LLM Team UI's OPENROUTER_API_KEY so both systems
// draw from one quota.
let k = v1::openrouter::resolve_openrouter_key();
if k.is_some() {
tracing::info!("v1: OpenRouter key loaded — /v1/chat provider=openrouter enabled");
} else {
tracing::warn!("v1: no OpenRouter key — openrouter rescue rung will 503");
}
k
},
gemini_key: {
// Phase 40 provider. GEMINI_API_KEY in env or .env.
let k = v1::gemini::resolve_gemini_key();
if k.is_some() {
tracing::info!("v1: Gemini key loaded — /v1/chat provider=gemini enabled");
} else {
tracing::debug!("v1: no Gemini key — provider=gemini will 503");
}
k
},
claude_key: {
// Phase 40 provider. ANTHROPIC_API_KEY in env or .env.
let k = v1::claude::resolve_claude_key();
if k.is_some() {
tracing::info!("v1: Claude key loaded — /v1/chat provider=claude enabled");
} else {
tracing::debug!("v1: no Claude key — provider=claude will 503");
}
k
},
// Phase 40 early deliverable — Langfuse trace emitter.
// Defaults match mcp-server/tracing.ts conventions so
// gateway traces land in the same staffing project.
langfuse: {
let c = v1::langfuse_trace::LangfuseClient::from_env_or_defaults();
if c.is_some() {
tracing::info!("v1: Langfuse tracing enabled — /v1/chat calls will appear at localhost:3001");
} else {
tracing::warn!("v1: Langfuse keys missing — /v1/chat calls will not be traced");
}
c
},
}));
// Auth middleware (if enabled) — P5-001 fix 2026-04-23:
// previously only inserted the ApiKey as an extension and never layered
// the middleware, so auth.enabled=true enforced nothing. Now wraps the
// router with from_fn_with_state, which calls api_key_auth on every
// request. /health is exempted inside the middleware (LB probes).
if config.auth.enabled {
if let Some(ref key) = config.auth.api_key {
tracing::info!("API key auth enabled — enforcing on all routes except /health");
let api_key = auth::ApiKey(key.clone());
app = app.layer(axum::middleware::from_fn_with_state(
api_key,
auth::api_key_auth,
));
} else {
tracing::warn!("auth enabled but no api_key set — all requests allowed");
}
}
app = app
.layer(DefaultBodyLimit::max(512 * 1024 * 1024)) // 512MB — supports 500K worker CSV
.layer(CorsLayer::new()
.allow_origin(Any)
.allow_methods(Any)
.allow_headers(Any))
.layer(TraceLayer::new_for_http());
// File watcher — auto-ingest from ./inbox
let _watcher = ingestd::watcher::spawn_watcher(
ingestd::watcher::WatchConfig::default(),
store.clone(),
registry.clone(),
);
// Start gRPC server on port+1
let grpc_port = config.gateway.port + 1;
let catalog_grpc = catalogd::grpc::CatalogGrpc::new(registry);
let grpc_addr = format!("{}:{}", config.gateway.host, grpc_port).parse().unwrap();
tokio::spawn(async move {
tracing::info!("gRPC server listening on {grpc_addr}");
tonic::transport::Server::builder()
.add_service(CatalogServiceServer::new(catalog_grpc))
.serve(grpc_addr)
.await
.expect("gRPC server failed");
});
// Start HTTP server
let http_addr = format!("{}:{}", config.gateway.host, config.gateway.port);
tracing::info!("HTTP gateway listening on {http_addr}");
let listener = tokio::net::TcpListener::bind(&http_addr).await.unwrap();
axum::serve(listener, app).await.unwrap();
}
async fn health() -> &'static str {
"lakehouse ok"
}
+37
View File
@@ -0,0 +1,37 @@
use opentelemetry::trace::TracerProvider;
use opentelemetry_sdk::trace::SdkTracerProvider;
use tracing_opentelemetry::OpenTelemetryLayer;
use tracing_subscriber::{EnvFilter, Registry, fmt, layer::SubscriberExt, util::SubscriberInitExt};
/// Initialize tracing with optional OpenTelemetry export.
pub fn init_tracing(service_name: &str, exporter: &str) {
let env_filter = EnvFilter::try_from_default_env()
.unwrap_or_else(|_| EnvFilter::new("info"));
match exporter {
"stdout" => {
// OTel traces to stdout + regular fmt layer
let provider = SdkTracerProvider::builder()
.with_simple_exporter(opentelemetry_stdout::SpanExporter::default())
.build();
let tracer = provider.tracer(service_name.to_string());
Registry::default()
.with(env_filter)
.with(fmt::layer())
.with(OpenTelemetryLayer::new(tracer))
.init();
tracing::info!("tracing initialized: fmt + opentelemetry-stdout");
}
_ => {
// Just fmt layer, no OTel
Registry::default()
.with(env_filter)
.with(fmt::layer())
.init();
tracing::info!("tracing initialized: fmt only");
}
}
}
+53
View File
@@ -0,0 +1,53 @@
pub mod registry;
pub mod service;
use queryd::context::QueryEngine;
use arrow::json::writer::{JsonArray, Writer as JsonWriter};
use std::sync::Arc;
use truth::TruthStore;
/// State for the tool system.
#[derive(Clone)]
pub struct ToolState {
pub registry: registry::ToolRegistry,
pub query_fn: QueryExecutor,
/// SQL guard (shared with queryd). Mirrors the queryd /sql truth
/// gate from P42-002 (9cc0ceb) — tools also execute model-
/// originated SQL, need the same destructive-verb block.
pub truth: Arc<TruthStore>,
}
/// Wraps QueryEngine to provide a simple execute interface for tools.
#[derive(Clone)]
pub struct QueryExecutor {
engine: QueryEngine,
}
impl QueryExecutor {
pub fn new(engine: QueryEngine) -> Self {
Self { engine }
}
/// Execute SQL and return (rows as JSON, row count).
pub async fn execute(&self, sql: &str) -> Result<(serde_json::Value, usize), String> {
let batches = self.engine.query(sql).await?;
if batches.is_empty() {
return Ok((serde_json::Value::Array(vec![]), 0));
}
let mut buf = Vec::new();
let mut writer = JsonWriter::<_, JsonArray>::new(&mut buf);
for batch in &batches {
writer.write(batch).map_err(|e| format!("JSON write: {e}"))?;
}
writer.finish().map_err(|e| format!("JSON finish: {e}"))?;
drop(writer);
let rows: serde_json::Value = serde_json::from_slice(&buf)
.map_err(|e| format!("JSON parse: {e}"))?;
let count = rows.as_array().map(|a| a.len()).unwrap_or(0);
Ok((rows, count))
}
}
+264
View File
@@ -0,0 +1,264 @@
/// Tool Registry: named, governed business actions for AI agents.
/// Instead of raw SQL, agents call validated tools with audit trails.
///
/// Each tool has:
/// - Name and description (for LLM tool-use)
/// - Parameter schema (validated before execution)
/// - Permission level (read / write / admin)
/// - Audit logging (every invocation recorded)
/// - Rate limiting (per agent)
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
/// Permission level for a tool.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Permission {
Read, // auto-approved, no side effects
Write, // modifies data, logged
Admin, // destructive, requires confirmation
}
/// Parameter definition for a tool.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ParamDef {
pub name: String,
pub param_type: String, // "string", "integer", "boolean", "float"
pub required: bool,
pub description: String,
pub default: Option<serde_json::Value>,
}
/// Tool definition — what agents see and can call.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolDef {
pub name: String,
pub description: String,
pub permission: Permission,
pub parameters: Vec<ParamDef>,
pub returns: String, // description of return value
pub sql_template: String, // SQL with {param} placeholders
pub category: String, // "candidates", "placements", "analytics"
}
/// Audit log entry for a tool invocation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolInvocation {
pub id: String,
pub tool_name: String,
pub agent: String,
pub params: serde_json::Value,
pub permission: Permission,
pub timestamp: DateTime<Utc>,
pub success: bool,
pub error: Option<String>,
pub rows_returned: Option<usize>,
}
/// The registry — holds tool definitions and audit log.
#[derive(Clone)]
pub struct ToolRegistry {
tools: Arc<RwLock<HashMap<String, ToolDef>>>,
audit_log: Arc<RwLock<Vec<ToolInvocation>>>,
}
impl ToolRegistry {
/// Build an empty registry. Callers in an async context should follow
/// this with `.register_defaults().await` if they want the built-in
/// staffing tools pre-installed — main.rs does exactly that.
pub fn new() -> Self {
Self {
tools: Arc::new(RwLock::new(HashMap::new())),
audit_log: Arc::new(RwLock::new(Vec::new())),
}
}
/// Register default staffing tools.
pub async fn register_defaults(&self) {
let tools = vec![
ToolDef {
name: "search_candidates".into(),
description: "Search for candidates by skills, city, state, availability, and experience. Returns matching candidates with contact info.".into(),
permission: Permission::Read,
parameters: vec![
ParamDef { name: "skills".into(), param_type: "string".into(), required: false, description: "Comma-separated skills to match (uses LIKE)".into(), default: None },
ParamDef { name: "city".into(), param_type: "string".into(), required: false, description: "City name".into(), default: None },
ParamDef { name: "state".into(), param_type: "string".into(), required: false, description: "State abbreviation".into(), default: None },
ParamDef { name: "min_years".into(), param_type: "integer".into(), required: false, description: "Minimum years of experience".into(), default: Some(serde_json::json!(0)) },
ParamDef { name: "status".into(), param_type: "string".into(), required: false, description: "Candidate status (active, inactive, placed)".into(), default: Some(serde_json::json!("active")) },
ParamDef { name: "limit".into(), param_type: "integer".into(), required: false, description: "Max results".into(), default: Some(serde_json::json!(20)) },
],
returns: "List of candidates with id, name, phone, email, skills, experience".into(),
sql_template: "SELECT candidate_id, first_name, last_name, phone, email, city, state, skills, years_experience, status FROM candidates WHERE 1=1 {skills_filter} {city_filter} {state_filter} {years_filter} {status_filter} ORDER BY years_experience DESC LIMIT {limit}".into(),
category: "candidates".into(),
},
ToolDef {
name: "get_candidate".into(),
description: "Get full details for a specific candidate by ID.".into(),
permission: Permission::Read,
parameters: vec![
ParamDef { name: "candidate_id".into(), param_type: "string".into(), required: true, description: "Candidate ID (e.g. CAND-000001)".into(), default: None },
],
returns: "Full candidate record".into(),
sql_template: "SELECT * FROM candidates WHERE candidate_id = '{candidate_id}'".into(),
category: "candidates".into(),
},
ToolDef {
name: "revenue_by_client".into(),
description: "Show total billed revenue, pay costs, and gross profit by client. Filter by date range.".into(),
permission: Permission::Read,
parameters: vec![
ParamDef { name: "limit".into(), param_type: "integer".into(), required: false, description: "Top N clients".into(), default: Some(serde_json::json!(10)) },
],
returns: "Client name, total billed, total paid, gross profit, timesheet count".into(),
sql_template: "SELECT c.company_name, COUNT(*) as timesheets, ROUND(SUM(t.bill_total),2) as total_billed, ROUND(SUM(t.pay_total),2) as total_paid, ROUND(SUM(t.bill_total) - SUM(t.pay_total),2) as gross_profit FROM timesheets t JOIN clients c ON t.client_id = c.client_id WHERE t.approved = true GROUP BY c.company_name ORDER BY total_billed DESC LIMIT {limit}".into(),
category: "analytics".into(),
},
ToolDef {
name: "recruiter_performance".into(),
description: "Show recruiter performance: placements, unique candidates, and total revenue generated.".into(),
permission: Permission::Read,
parameters: vec![
ParamDef { name: "limit".into(), param_type: "integer".into(), required: false, description: "Top N recruiters".into(), default: Some(serde_json::json!(10)) },
],
returns: "Recruiter name, placement count, unique candidates, total revenue".into(),
sql_template: "SELECT p.recruiter, COUNT(DISTINCT p.placement_id) as placements, COUNT(DISTINCT p.candidate_id) as unique_candidates, ROUND(SUM(t.bill_total),2) as total_revenue FROM placements p JOIN timesheets t ON p.placement_id = t.placement_id GROUP BY p.recruiter ORDER BY total_revenue DESC LIMIT {limit}".into(),
category: "analytics".into(),
},
ToolDef {
name: "cold_leads".into(),
description: "Find candidates who were called multiple times but never placed — potential lost opportunities.".into(),
permission: Permission::Read,
parameters: vec![
ParamDef { name: "min_calls".into(), param_type: "integer".into(), required: false, description: "Minimum call count".into(), default: Some(serde_json::json!(5)) },
ParamDef { name: "limit".into(), param_type: "integer".into(), required: false, description: "Max results".into(), default: Some(serde_json::json!(20)) },
],
returns: "Candidates with call counts who were never placed".into(),
sql_template: "SELECT c.candidate_id, c.first_name, c.last_name, c.phone, c.vertical, cl.calls FROM candidates c JOIN (SELECT candidate_id, COUNT(*) as calls FROM call_log GROUP BY candidate_id HAVING COUNT(*) >= {min_calls}) cl ON c.candidate_id = cl.candidate_id WHERE c.candidate_id NOT IN (SELECT DISTINCT candidate_id FROM placements) ORDER BY cl.calls DESC LIMIT {limit}".into(),
category: "analytics".into(),
},
ToolDef {
name: "open_jobs".into(),
description: "List open job orders with client, title, rates, and location.".into(),
permission: Permission::Read,
parameters: vec![
ParamDef { name: "vertical".into(), param_type: "string".into(), required: false, description: "Filter by vertical (IT, Healthcare, etc)".into(), default: None },
ParamDef { name: "city".into(), param_type: "string".into(), required: false, description: "Filter by city".into(), default: None },
ParamDef { name: "limit".into(), param_type: "integer".into(), required: false, description: "Max results".into(), default: Some(serde_json::json!(20)) },
],
returns: "Open job orders with details".into(),
sql_template: "SELECT j.job_order_id, c.company_name, j.title, j.vertical, j.city, j.state, j.bill_rate, j.pay_rate, j.status FROM job_orders j JOIN clients c ON j.client_id = c.client_id WHERE j.status = 'open' {vertical_filter} {city_filter} ORDER BY j.bill_rate DESC LIMIT {limit}".into(),
category: "jobs".into(),
},
];
let mut reg = self.tools.write().await;
for tool in tools {
reg.insert(tool.name.clone(), tool);
}
}
/// Get tool definition (for LLM tool-use schema).
pub async fn get_tool(&self, name: &str) -> Option<ToolDef> {
self.tools.read().await.get(name).cloned()
}
/// List all tools (for LLM tool discovery).
pub async fn list_tools(&self) -> Vec<ToolDef> {
self.tools.read().await.values().cloned().collect()
}
/// Build SQL from tool parameters. Validates and sanitizes.
pub fn build_sql(tool: &ToolDef, params: &serde_json::Value) -> Result<String, String> {
let mut sql = tool.sql_template.clone();
let params = params.as_object().ok_or("params must be an object")?;
// Replace named parameters
for param_def in &tool.parameters {
let value = params.get(&param_def.name)
.or(param_def.default.as_ref());
match param_def.name.as_str() {
// Handle filter parameters (conditional WHERE clauses)
"skills" => {
if let Some(v) = value.and_then(|v| v.as_str()).filter(|s| !s.is_empty()) {
let filters: Vec<String> = v.split(',')
.map(|s| format!("skills LIKE '%{}%'", s.trim().replace('\'', "''")))
.collect();
sql = sql.replace("{skills_filter}", &format!("AND ({})", filters.join(" OR ")));
} else {
sql = sql.replace("{skills_filter}", "");
}
}
"city" => {
if let Some(v) = value.and_then(|v| v.as_str()).filter(|s| !s.is_empty()) {
sql = sql.replace("{city_filter}", &format!("AND city = '{}'", v.replace('\'', "''")));
} else {
sql = sql.replace("{city_filter}", "");
}
}
"state" => {
if let Some(v) = value.and_then(|v| v.as_str()).filter(|s| !s.is_empty()) {
sql = sql.replace("{state_filter}", &format!("AND state = '{}'", v.replace('\'', "''")));
} else {
sql = sql.replace("{state_filter}", "");
}
}
"vertical" => {
if let Some(v) = value.and_then(|v| v.as_str()).filter(|s| !s.is_empty()) {
sql = sql.replace("{vertical_filter}", &format!("AND j.vertical = '{}'", v.replace('\'', "''")));
} else {
sql = sql.replace("{vertical_filter}", "");
}
}
"status" => {
if let Some(v) = value.and_then(|v| v.as_str()).filter(|s| !s.is_empty()) {
sql = sql.replace("{status_filter}", &format!("AND status = '{}'", v.replace('\'', "''")));
} else {
sql = sql.replace("{status_filter}", "");
}
}
"min_years" => {
if let Some(v) = value.and_then(|v| v.as_i64()) {
sql = sql.replace("{years_filter}", &format!("AND years_experience >= {v}"));
} else {
sql = sql.replace("{years_filter}", "");
}
}
// Direct substitution for simple params
_ => {
let placeholder = format!("{{{}}}", param_def.name);
if sql.contains(&placeholder) {
let val_str = match value {
Some(serde_json::Value::String(s)) => s.replace('\'', "''"),
Some(serde_json::Value::Number(n)) => n.to_string(),
Some(serde_json::Value::Bool(b)) => b.to_string(),
Some(v) => v.to_string(),
None if param_def.required => return Err(format!("missing required param: {}", param_def.name)),
None => continue,
};
sql = sql.replace(&placeholder, &val_str);
}
}
}
}
Ok(sql)
}
/// Log a tool invocation.
pub async fn log_invocation(&self, inv: ToolInvocation) {
self.audit_log.write().await.push(inv);
}
/// Get recent audit log.
pub async fn recent_audit(&self, limit: usize) -> Vec<ToolInvocation> {
let log = self.audit_log.read().await;
let start = log.len().saturating_sub(limit);
log[start..].iter().rev().cloned().collect()
}
}
+174
View File
@@ -0,0 +1,174 @@
use axum::{
Json, Router,
extract::{Path, Query, State},
http::StatusCode,
response::IntoResponse,
routing::{get, post},
};
use serde::Deserialize;
use super::registry::{ToolInvocation, ToolRegistry};
use crate::tools::ToolState;
pub fn router(state: ToolState) -> Router {
Router::new()
.route("/", get(list_tools))
.route("/{name}", get(get_tool))
.route("/{name}/call", post(call_tool))
.route("/audit", get(audit_log))
.with_state(state)
}
/// List all available tools (for LLM tool-use discovery).
async fn list_tools(State(state): State<ToolState>) -> impl IntoResponse {
let tools = state.registry.list_tools().await;
// Return in MCP-compatible format
let tool_list: Vec<serde_json::Value> = tools.iter().map(|t| {
serde_json::json!({
"name": t.name,
"description": t.description,
"permission": t.permission,
"category": t.category,
"parameters": t.parameters.iter().map(|p| {
serde_json::json!({
"name": p.name,
"type": p.param_type,
"required": p.required,
"description": p.description,
"default": p.default,
})
}).collect::<Vec<_>>(),
})
}).collect();
Json(tool_list)
}
/// Get a specific tool definition.
async fn get_tool(
State(state): State<ToolState>,
Path(name): Path<String>,
) -> impl IntoResponse {
match state.registry.get_tool(&name).await {
Some(tool) => Ok(Json(tool)),
None => Err((StatusCode::NOT_FOUND, format!("tool not found: {name}"))),
}
}
/// Call a tool — validate params, execute SQL, log invocation.
#[derive(Deserialize)]
struct CallRequest {
params: serde_json::Value,
agent: String,
}
async fn call_tool(
State(state): State<ToolState>,
Path(name): Path<String>,
Json(req): Json<CallRequest>,
) -> impl IntoResponse {
let tool = match state.registry.get_tool(&name).await {
Some(t) => t,
None => return Err((StatusCode::NOT_FOUND, format!("tool not found: {name}"))),
};
tracing::info!("tool call: {} by agent '{}' with {:?}", name, req.agent, req.params);
// Build SQL from params
let sql = match ToolRegistry::build_sql(&tool, &req.params) {
Ok(s) => s,
Err(e) => {
state.registry.log_invocation(ToolInvocation {
id: format!("inv-{}", chrono::Utc::now().timestamp_millis()),
tool_name: name.clone(),
agent: req.agent.clone(),
params: req.params.clone(),
permission: tool.permission.clone(),
timestamp: chrono::Utc::now(),
success: false,
error: Some(e.clone()),
rows_returned: None,
}).await;
return Err((StatusCode::BAD_REQUEST, e));
}
};
// Truth gate — same contract as queryd /sql (P42-002). Rejects
// destructive verbs + empty SQL. Scrum iter 11 CF-1 + CF-2 on this
// file: tools executed model-provided SQL parameters without any
// validation. Close the gap here so the parallel surface has the
// same safety floor as queryd.
let ctx = serde_json::json!({ "sql": sql });
for outcome in state.truth.evaluate("sql_query", &ctx) {
if outcome.passed {
if let truth::RuleAction::Reject { message } | truth::RuleAction::Block { message } = &outcome.action {
tracing::warn!("tool {name}: SQL blocked by truth gate ({}): {message}", outcome.rule_id);
state.registry.log_invocation(ToolInvocation {
id: format!("inv-{}", chrono::Utc::now().timestamp_millis()),
tool_name: name.clone(),
agent: req.agent.clone(),
params: req.params.clone(),
permission: tool.permission.clone(),
timestamp: chrono::Utc::now(),
success: false,
error: Some(format!("truth gate: {message}")),
rows_returned: None,
}).await;
return Err((StatusCode::FORBIDDEN, message.clone()));
}
}
}
// Execute via query engine
let result = state.query_fn.execute(&sql).await;
match result {
Ok((rows, row_count)) => {
state.registry.log_invocation(ToolInvocation {
id: format!("inv-{}", chrono::Utc::now().timestamp_millis()),
tool_name: name.clone(),
agent: req.agent,
params: req.params,
permission: tool.permission,
timestamp: chrono::Utc::now(),
success: true,
error: None,
rows_returned: Some(row_count),
}).await;
Ok(Json(serde_json::json!({
"tool": name,
"rows": rows,
"row_count": row_count,
"sql": sql,
})))
}
Err(e) => {
state.registry.log_invocation(ToolInvocation {
id: format!("inv-{}", chrono::Utc::now().timestamp_millis()),
tool_name: name,
agent: req.agent,
params: req.params,
permission: tool.permission,
timestamp: chrono::Utc::now(),
success: false,
error: Some(e.clone()),
rows_returned: None,
}).await;
Err((StatusCode::INTERNAL_SERVER_ERROR, e))
}
}
}
#[derive(Deserialize)]
struct AuditQuery {
limit: Option<usize>,
}
async fn audit_log(
State(state): State<ToolState>,
Query(q): Query<AuditQuery>,
) -> impl IntoResponse {
let log = state.registry.recent_audit(q.limit.unwrap_or(50)).await;
Json(log)
}
+222
View File
@@ -0,0 +1,222 @@
//! Claude (Anthropic) adapter.
//!
//! POST `https://api.anthropic.com/v1/messages`. Auth via `x-api-key`
//! header (not bearer) + required `anthropic-version` header. Payload
//! is NOT OpenAI-compatible — response text lives at
//! `content[0].text`. Phase 40 deliverable. System prompts travel in
//! a top-level `system` field, separate from the `messages` array.
use std::time::Duration;
use serde::{Deserialize, Serialize};
use super::{ChatRequest, ChatResponse, Choice, Message, UsageBlock};
const CLAUDE_BASE_URL: &str = "https://api.anthropic.com/v1";
const CLAUDE_API_VERSION: &str = "2023-06-01";
const CLAUDE_TIMEOUT_SECS: u64 = 180;
pub fn resolve_claude_key() -> Option<String> {
if let Ok(k) = std::env::var("ANTHROPIC_API_KEY") {
if !k.trim().is_empty() { return Some(k.trim().to_string()); }
}
for path in ["/home/profit/.env", "/root/.env"] {
if let Ok(raw) = std::fs::read_to_string(path) {
for line in raw.lines() {
if let Some(rest) = line.strip_prefix("ANTHROPIC_API_KEY=") {
let k = rest.trim().trim_matches('"').trim_matches('\'');
if !k.is_empty() { return Some(k.to_string()); }
}
}
}
}
None
}
pub async fn chat(
key: &str,
req: &ChatRequest,
) -> Result<ChatResponse, String> {
// Strip the "claude/" prefix if the caller used the namespaced form.
let model = req.model.strip_prefix("claude/").unwrap_or(&req.model).to_string();
// Anthropic carries system prompts outside the messages array.
// Concatenate any system-role messages into a single system string;
// keep user + assistant messages in `messages`.
let mut system_parts: Vec<String> = Vec::new();
let mut msgs: Vec<AnMessage> = Vec::new();
for m in &req.messages {
if m.role == "system" {
system_parts.push(m.content.clone());
} else {
// Anthropic expects strictly "user" or "assistant"; anything
// else we normalize to "user".
let role = if m.role == "assistant" { "assistant" } else { "user" };
msgs.push(AnMessage { role: role.to_string(), content: m.content.clone() });
}
}
let system = if system_parts.is_empty() {
None
} else {
Some(system_parts.join("\n\n"))
};
let body = AnChatBody {
model: model.clone(),
messages: msgs,
max_tokens: req.max_tokens.unwrap_or(800),
temperature: req.temperature.unwrap_or(0.3),
system,
};
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(CLAUDE_TIMEOUT_SECS))
.build()
.map_err(|e| format!("build client: {e}"))?;
let t0 = std::time::Instant::now();
let resp = client
.post(format!("{}/messages", CLAUDE_BASE_URL))
.header("x-api-key", key)
.header("anthropic-version", CLAUDE_API_VERSION)
.json(&body)
.send()
.await
.map_err(|e| format!("api.anthropic.com unreachable: {e}"))?;
let status = resp.status();
if !status.is_success() {
let body = resp.text().await.unwrap_or_else(|_| "?".into());
return Err(format!("claude {}: {}", status, body));
}
let parsed: AnChatResponse = resp.json().await
.map_err(|e| format!("invalid claude response: {e}"))?;
let latency_ms = t0.elapsed().as_millis();
let text = parsed.content.into_iter()
.find(|b| b.block_type == "text")
.map(|b| b.text)
.unwrap_or_default();
let prompt_tokens = parsed.usage.as_ref().map(|u| u.input_tokens).unwrap_or_else(|| {
let chars: usize = req.messages.iter().map(|m| m.content.chars().count()).sum();
((chars + 3) / 4) as u32
});
let completion_tokens = parsed.usage.as_ref().map(|u| u.output_tokens).unwrap_or_else(|| {
((text.chars().count() + 3) / 4) as u32
});
tracing::info!(
target: "v1.chat",
provider = "claude",
model = %model,
prompt_tokens,
completion_tokens,
latency_ms = latency_ms as u64,
"claude chat completed",
);
Ok(ChatResponse {
id: format!("chatcmpl-{}", chrono::Utc::now().timestamp_nanos_opt().unwrap_or(0)),
object: "chat.completion",
created: chrono::Utc::now().timestamp(),
model,
choices: vec![Choice {
index: 0,
message: Message { role: "assistant".into(), content: text },
finish_reason: parsed.stop_reason.unwrap_or_else(|| "stop".into()),
}],
usage: UsageBlock {
prompt_tokens,
completion_tokens,
total_tokens: prompt_tokens + completion_tokens,
},
})
}
// -- Anthropic Messages API wire shapes --
#[derive(Serialize)]
struct AnChatBody {
model: String,
messages: Vec<AnMessage>,
max_tokens: u32,
temperature: f64,
#[serde(skip_serializing_if = "Option::is_none")]
system: Option<String>,
}
#[derive(Serialize)]
struct AnMessage { role: String, content: String }
#[derive(Deserialize)]
struct AnChatResponse {
content: Vec<AnContentBlock>,
#[serde(default, rename = "stop_reason")]
stop_reason: Option<String>,
#[serde(default)]
usage: Option<AnUsage>,
}
#[derive(Deserialize)]
struct AnContentBlock {
#[serde(rename = "type")]
block_type: String,
#[serde(default)]
text: String,
}
#[derive(Deserialize)]
struct AnUsage { input_tokens: u32, output_tokens: u32 }
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn resolve_claude_key_does_not_panic() {
let _ = resolve_claude_key();
}
#[test]
fn chat_body_serializes_with_separate_system() {
let body = AnChatBody {
model: "claude-3-5-sonnet-latest".into(),
messages: vec![
AnMessage { role: "user".into(), content: "hi".into() },
],
max_tokens: 800,
temperature: 0.3,
system: Some("You are helpful.".into()),
};
let json = serde_json::to_string(&body).unwrap();
assert!(json.contains("\"system\":\"You are helpful.\""));
assert!(json.contains("\"messages\""));
assert!(json.contains("\"max_tokens\":800"));
}
#[test]
fn body_omits_system_when_none() {
let body = AnChatBody {
model: "claude-3-5-sonnet-latest".into(),
messages: vec![AnMessage { role: "user".into(), content: "hi".into() }],
max_tokens: 800,
temperature: 0.3,
system: None,
};
let json = serde_json::to_string(&body).unwrap();
assert!(!json.contains("\"system\""), "system field should be skipped when None: {json}");
}
#[test]
fn model_prefix_strip_preserves_bare_names() {
let cases = [
("claude/claude-3-5-sonnet-latest", "claude-3-5-sonnet-latest"),
("claude-3-5-sonnet-latest", "claude-3-5-sonnet-latest"),
];
for (input, expected) in cases {
let out = input.strip_prefix("claude/").unwrap_or(input);
assert_eq!(out, expected);
}
}
}
+230
View File
@@ -0,0 +1,230 @@
//! Gemini adapter — Google's Generative Language API.
//!
//! POST `https://generativelanguage.googleapis.com/v1beta/models/
//! {model}:generateContent?key=<API_KEY>`. Auth via query-string key
//! (not bearer). Payload shape is NOT OpenAI-compatible — we map
//! messages → contents + parts, extract response from `candidates[0]
//! .content.parts[0].text`. Phase 40 deliverable; gate: `/v1/chat`
//! with a prefixed or explicit gemini model returns normally.
use std::time::Duration;
use serde::{Deserialize, Serialize};
use super::{ChatRequest, ChatResponse, Choice, Message, UsageBlock};
const GEMINI_BASE_URL: &str = "https://generativelanguage.googleapis.com/v1beta";
const GEMINI_TIMEOUT_SECS: u64 = 180;
pub fn resolve_gemini_key() -> Option<String> {
if let Ok(k) = std::env::var("GEMINI_API_KEY") {
if !k.trim().is_empty() { return Some(k.trim().to_string()); }
}
for path in ["/home/profit/.env", "/root/.env"] {
if let Ok(raw) = std::fs::read_to_string(path) {
for line in raw.lines() {
if let Some(rest) = line.strip_prefix("GEMINI_API_KEY=") {
let k = rest.trim().trim_matches('"').trim_matches('\'');
if !k.is_empty() { return Some(k.to_string()); }
}
}
}
}
None
}
pub async fn chat(
key: &str,
req: &ChatRequest,
) -> Result<ChatResponse, String> {
// Strip the "gemini/" prefix if the caller used the namespaced form.
let model = req.model.strip_prefix("gemini/").unwrap_or(&req.model).to_string();
// Gemini splits system prompt from conversation differently.
// Simplest working mapping: concatenate any system messages at the
// top of a single user turn, then append user/assistant turns as
// separate contents entries. Covers the common single-turn case
// the scrum pipeline uses.
let mut contents: Vec<GmContent> = Vec::new();
for m in &req.messages {
let role = match m.role.as_str() {
"system" | "user" => "user",
_ => "model",
};
contents.push(GmContent {
role: role.to_string(),
parts: vec![GmPart { text: m.content.clone() }],
});
}
let body = GmChatBody {
contents,
generation_config: GmGenerationConfig {
temperature: req.temperature.unwrap_or(0.3),
max_output_tokens: req.max_tokens.unwrap_or(800),
},
};
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(GEMINI_TIMEOUT_SECS))
.build()
.map_err(|e| format!("build client: {e}"))?;
let url = format!("{}/models/{}:generateContent?key={}", GEMINI_BASE_URL, model, key);
let t0 = std::time::Instant::now();
let resp = client
.post(&url)
.json(&body)
.send()
.await
.map_err(|e| format!("generativelanguage.googleapis.com unreachable: {e}"))?;
let status = resp.status();
if !status.is_success() {
let body = resp.text().await.unwrap_or_else(|_| "?".into());
return Err(format!("gemini {}: {}", status, body));
}
let parsed: GmChatResponse = resp.json().await
.map_err(|e| format!("invalid gemini response: {e}"))?;
let latency_ms = t0.elapsed().as_millis();
let candidate = parsed.candidates.into_iter().next()
.ok_or_else(|| "gemini returned no candidates".to_string())?;
let text = candidate.content.parts.into_iter()
.next()
.map(|p| p.text)
.unwrap_or_default();
let prompt_tokens = parsed.usage_metadata.as_ref()
.map(|u| u.prompt_token_count)
.unwrap_or_else(|| {
let chars: usize = req.messages.iter().map(|m| m.content.chars().count()).sum();
((chars + 3) / 4) as u32
});
let completion_tokens = parsed.usage_metadata.as_ref()
.map(|u| u.candidates_token_count)
.unwrap_or_else(|| ((text.chars().count() + 3) / 4) as u32);
tracing::info!(
target: "v1.chat",
provider = "gemini",
model = %model,
prompt_tokens,
completion_tokens,
latency_ms = latency_ms as u64,
"gemini chat completed",
);
Ok(ChatResponse {
id: format!("chatcmpl-{}", chrono::Utc::now().timestamp_nanos_opt().unwrap_or(0)),
object: "chat.completion",
created: chrono::Utc::now().timestamp(),
model,
choices: vec![Choice {
index: 0,
message: Message { role: "assistant".into(), content: text },
finish_reason: candidate.finish_reason.unwrap_or_else(|| "stop".into()),
}],
usage: UsageBlock {
prompt_tokens,
completion_tokens,
total_tokens: prompt_tokens + completion_tokens,
},
})
}
// -- Gemini wire shapes --
#[derive(Serialize)]
struct GmChatBody {
contents: Vec<GmContent>,
#[serde(rename = "generationConfig")]
generation_config: GmGenerationConfig,
}
#[derive(Serialize)]
struct GmContent {
role: String,
parts: Vec<GmPart>,
}
#[derive(Serialize)]
struct GmPart { text: String }
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct GmGenerationConfig {
temperature: f64,
max_output_tokens: u32,
}
#[derive(Deserialize)]
struct GmChatResponse {
candidates: Vec<GmCandidate>,
#[serde(default, rename = "usageMetadata")]
usage_metadata: Option<GmUsage>,
}
#[derive(Deserialize)]
struct GmCandidate {
content: GmContentResp,
#[serde(default, rename = "finishReason")]
finish_reason: Option<String>,
}
#[derive(Deserialize)]
struct GmContentResp { parts: Vec<GmPartResp> }
#[derive(Deserialize)]
struct GmPartResp { #[serde(default)] text: String }
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
struct GmUsage {
prompt_token_count: u32,
candidates_token_count: u32,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn resolve_gemini_key_does_not_panic() {
let _ = resolve_gemini_key();
}
#[test]
fn chat_body_serializes_to_gemini_shape() {
let body = GmChatBody {
contents: vec![
GmContent {
role: "user".into(),
parts: vec![GmPart { text: "hello".into() }],
},
],
generation_config: GmGenerationConfig {
temperature: 0.3,
max_output_tokens: 800,
},
};
let json = serde_json::to_string(&body).unwrap();
assert!(json.contains("\"contents\""));
assert!(json.contains("\"parts\""));
// camelCase per Gemini API
assert!(json.contains("\"generationConfig\""));
assert!(json.contains("\"maxOutputTokens\":800"));
}
#[test]
fn model_prefix_strip_preserves_bare_names() {
let cases = [
("gemini/gemini-2.0-flash", "gemini-2.0-flash"),
("gemini-2.0-flash", "gemini-2.0-flash"),
];
for (input, expected) in cases {
let out = input.strip_prefix("gemini/").unwrap_or(input);
assert_eq!(out, expected);
}
}
}
+226
View File
@@ -0,0 +1,226 @@
//! Phase 40 (early deliverable) — Langfuse tracing from the Rust gateway.
//!
//! Langfuse is already running at `http://localhost:3001` with the
//! staffing project — `mcp-server/tracing.ts` has been emitting traces
//! for months. This file bridges the gap: every `/v1/chat` call now
//! appears in the same Langfuse project, alongside traces from the
//! Bun side.
//!
//! Design:
//! - Fire-and-forget — emit spawns a tokio task. Response to client
//! never waits on Langfuse, never fails because of it. If Langfuse
//! is down or slow, we silently drop traces (and `tracing::warn!`).
//! - Two events per call — `trace-create` + `generation-create`. The
//! generation nests under the trace so Langfuse UI shows the whole
//! thing as one clickable row.
//! - No SDK — plain HTTP POST to `/api/public/ingestion` with Basic
//! auth. Keeps the Rust dependency graph small.
use serde::Serialize;
use std::sync::Arc;
use std::time::Duration;
use super::Message;
const INGESTION_PATH: &str = "/api/public/ingestion";
const EMIT_TIMEOUT_SECS: u64 = 5;
#[derive(Clone)]
pub struct LangfuseClient {
inner: Arc<Inner>,
}
struct Inner {
base_url: String,
public_key: String,
secret_key: String,
http: reqwest::Client,
}
impl LangfuseClient {
pub fn new(base_url: String, public_key: String, secret_key: String) -> Self {
let http = reqwest::Client::builder()
.timeout(Duration::from_secs(EMIT_TIMEOUT_SECS))
.build()
.expect("langfuse client init");
Self { inner: Arc::new(Inner { base_url, public_key, secret_key, http }) }
}
/// Resolve a LangfuseClient from the same conventions the TS side
/// already uses. Returns None if keys are missing — caller logs a
/// warning and omits tracing (non-fatal).
///
/// Sources checked in order:
/// 1. LANGFUSE_URL + LANGFUSE_PUBLIC_KEY + LANGFUSE_SECRET_KEY env
/// 2. /root/llm_team_config.json (if it carries langfuse keys)
/// 3. Hardcoded staffing-project defaults matching mcp-server
/// (pk-lf-staffing / sk-lf-staffing-secret / localhost:3001)
pub fn from_env_or_defaults() -> Option<Self> {
let base = std::env::var("LANGFUSE_URL")
.unwrap_or_else(|_| "http://localhost:3001".to_string());
let pk = std::env::var("LANGFUSE_PUBLIC_KEY")
.unwrap_or_else(|_| "pk-lf-staffing".to_string());
let sk = std::env::var("LANGFUSE_SECRET_KEY")
.unwrap_or_else(|_| "sk-lf-staffing-secret".to_string());
if pk.trim().is_empty() || sk.trim().is_empty() { return None; }
Some(Self::new(base, pk, sk))
}
/// Fire-and-forget emit. Never blocks, never errors upward.
pub fn emit_chat(&self, ev: ChatTrace) {
let this = self.clone();
tokio::spawn(async move {
if let Err(e) = this.emit_chat_inner(ev).await {
tracing::warn!(target: "v1.langfuse", "trace drop: {e}");
}
});
}
async fn emit_chat_inner(&self, ev: ChatTrace) -> Result<(), String> {
let trace_id = uuid_v7_like();
let gen_id = uuid_v7_like();
let trace_ts = ev.start_time.clone();
let batch = IngestionBatch {
batch: vec![
IngestionEvent {
id: uuid_v7_like(),
timestamp: trace_ts.clone(),
kind: "trace-create",
body: serde_json::json!({
"id": trace_id,
"name": format!("v1.chat:{}", ev.provider),
"input": serde_json::json!({
"model": ev.model,
"messages": ev.input,
}),
"metadata": serde_json::json!({
"provider": ev.provider,
"think": ev.think,
}),
}),
},
IngestionEvent {
id: uuid_v7_like(),
timestamp: ev.end_time.clone(),
kind: "generation-create",
body: serde_json::json!({
"id": gen_id,
"traceId": trace_id,
"name": "chat",
"model": ev.model,
"modelParameters": serde_json::json!({
"temperature": ev.temperature,
"max_tokens": ev.max_tokens,
"think": ev.think,
}),
"input": ev.input,
"output": ev.output,
"usage": serde_json::json!({
"input": ev.prompt_tokens,
"output": ev.completion_tokens,
"total": ev.prompt_tokens + ev.completion_tokens,
"unit": "TOKENS",
}),
"startTime": ev.start_time,
"endTime": ev.end_time,
"metadata": serde_json::json!({
"provider": ev.provider,
"latency_ms": ev.latency_ms,
}),
}),
},
],
};
let url = format!("{}{}", self.inner.base_url.trim_end_matches('/'), INGESTION_PATH);
let resp = self.inner.http
.post(url)
.basic_auth(&self.inner.public_key, Some(&self.inner.secret_key))
.json(&batch)
.send()
.await
.map_err(|e| format!("POST failed: {e}"))?;
if !resp.status().is_success() {
return Err(format!("{}: {}", resp.status(), resp.text().await.unwrap_or_default()));
}
Ok(())
}
}
/// Everything the v1.chat handler collects for one completed call.
pub struct ChatTrace {
pub provider: String,
pub model: String,
pub input: Vec<Message>,
pub output: String,
pub prompt_tokens: u32,
pub completion_tokens: u32,
pub temperature: Option<f64>,
pub max_tokens: Option<u32>,
pub think: Option<bool>,
pub start_time: String,
pub end_time: String,
pub latency_ms: u64,
}
#[derive(Serialize)]
struct IngestionBatch {
batch: Vec<IngestionEvent>,
}
#[derive(Serialize)]
struct IngestionEvent {
id: String,
timestamp: String,
#[serde(rename = "type")]
kind: &'static str,
body: serde_json::Value,
}
/// UUIDv7-ish identifier — not strictly UUIDv7 but time-ordered and
/// unique enough for event deduplication. Avoids a uuid crate dep.
fn uuid_v7_like() -> String {
let ts = chrono::Utc::now().timestamp_nanos_opt().unwrap_or(0);
let rand = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.subsec_nanos())
.unwrap_or(0);
format!("{:016x}-{:08x}", ts, rand)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn from_env_or_defaults_returns_some_when_defaults_present() {
// Defaults are hardcoded non-empty — should always produce Some.
let c = LangfuseClient::from_env_or_defaults();
assert!(c.is_some());
}
#[test]
fn uuid_v7_like_produces_distinct_ids() {
let a = uuid_v7_like();
// Different nanosecond timestamps — but they might collide on
// back-to-back calls. Loop a bit to ensure uniqueness surface.
let b = uuid_v7_like();
let c = uuid_v7_like();
assert!(a != b || b != c, "three consecutive ids should have at least one distinct pair");
assert!(a.contains('-'));
}
#[test]
fn ingestion_batch_serializes_with_rename() {
let ev = IngestionEvent {
id: "e1".into(),
timestamp: "2026-04-22T00:00:00Z".into(),
kind: "trace-create",
body: serde_json::json!({"name": "t"}),
};
let json = serde_json::to_string(&IngestionBatch { batch: vec![ev] }).unwrap();
assert!(json.contains("\"type\":\"trace-create\""));
assert!(json.contains("\"batch\":["));
}
}
+440
View File
@@ -0,0 +1,440 @@
//! Phase 38 — Universal API skeleton (`/v1/*`).
//!
//! OpenAI-compatible shape on top of the existing aibridge → Ollama
//! path. This is the thin slice: single provider, stateless, no
//! streaming. Phase 39 replaces the direct Ollama call with a
//! `ProviderAdapter` trait dispatch; Phase 40 adds routing + fallback.
//!
//! The shape matches OpenAI's `/v1/chat/completions` closely enough
//! that clients using openai-compatible SDKs can point at us with the
//! URL swap alone. We keep the endpoint path `/v1/chat` (not
//! `/v1/chat/completions`) because our PRD declares the terser form;
//! adding the alias is one line in Phase 39 when it matters.
pub mod ollama;
pub mod ollama_cloud;
pub mod openrouter;
pub mod gemini;
pub mod claude;
pub mod langfuse_trace;
pub mod respond;
pub mod truth;
use axum::{
Router,
extract::State,
http::StatusCode,
response::IntoResponse,
routing::{get, post},
Json,
};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tokio::sync::RwLock;
#[derive(Clone)]
pub struct V1State {
pub ai_client: aibridge::client::AiClient,
pub usage: Arc<RwLock<Usage>>,
/// Ollama Cloud bearer token. Loaded at startup via
/// `ollama_cloud::resolve_cloud_key()`. None = cloud routes 503.
pub ollama_cloud_key: Option<String>,
/// OpenRouter bearer token — free-tier rescue rung. Loaded at
/// startup via `openrouter::resolve_openrouter_key()`. None means
/// provider="openrouter" calls 503 rather than attempt. Same key
/// sourcing as LLM Team UI so the two share one API quota.
pub openrouter_key: Option<String>,
/// Gemini API key (Google Generative Language). Loaded at startup
/// via `gemini::resolve_gemini_key()`. None = provider="gemini"
/// calls 503. Phase 40 deliverable.
pub gemini_key: Option<String>,
/// Anthropic Claude API key. Loaded at startup via
/// `claude::resolve_claude_key()`. None = provider="claude" calls
/// 503. Phase 40 deliverable.
pub claude_key: Option<String>,
/// Phase 40 early deliverable — Langfuse client. None = tracing
/// disabled (keys missing or container unreachable). Traces are
/// fire-and-forget: never block the response path.
pub langfuse: Option<langfuse_trace::LangfuseClient>,
}
#[derive(Default, Clone, Serialize)]
pub struct Usage {
pub requests: u64,
pub prompt_tokens: u64,
pub completion_tokens: u64,
pub total_tokens: u64,
#[serde(default)]
pub by_provider: std::collections::HashMap<String, ProviderUsage>,
}
#[derive(Default, Clone, Serialize)]
pub struct ProviderUsage {
pub requests: u64,
pub prompt_tokens: u64,
pub completion_tokens: u64,
pub total_tokens: u64,
}
pub fn router(state: V1State) -> Router {
Router::new()
.route("/chat", post(chat))
.route("/respond", post(respond::respond))
.route("/usage", get(usage))
.route("/sessions", get(sessions))
.route("/context", get(truth::context))
.with_state(state)
}
// -- Shared types (OpenAI-compatible) --
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct Message {
pub role: String,
pub content: String,
}
#[derive(Deserialize, Debug, Clone)]
pub struct ChatRequest {
pub model: String,
pub messages: Vec<Message>,
#[serde(default)]
pub temperature: Option<f64>,
#[serde(default)]
pub max_tokens: Option<u32>,
/// Accepted for shape-compat but ignored in the thin slice —
/// Phase 38 returns non-streaming even when the client asked for it.
/// Phase 39+ wires real streaming.
#[serde(default)]
pub stream: Option<bool>,
/// Non-OpenAI extension. Passes through to the provider's thinking
/// toggle. Default: **false** — hot-path discipline for thinking
/// models (qwen3.5, qwen3, gpt-oss) that otherwise burn the token
/// budget on hidden reasoning before visible output starts,
/// producing empty responses. Set true explicitly when calling an
/// overseer / reasoning-heavy path.
#[serde(default)]
pub think: Option<bool>,
/// Non-OpenAI extension. Selects the provider adapter. Accepted:
/// - None / "ollama" / "local" → local Ollama via sidecar (default)
/// - "ollama_cloud" / "cloud" → Ollama Cloud direct HTTPS
/// Phase 40 adds a routing engine that picks this automatically
/// from the model name; Phase 38/39 requires explicit selection.
#[serde(default)]
pub provider: Option<String>,
}
#[derive(Serialize)]
pub struct ChatResponse {
pub id: String,
pub object: &'static str,
pub created: i64,
pub model: String,
pub choices: Vec<Choice>,
pub usage: UsageBlock,
}
#[derive(Serialize)]
pub struct Choice {
pub index: u32,
pub message: Message,
pub finish_reason: String,
}
#[derive(Serialize, Deserialize, Clone)]
pub struct UsageBlock {
pub prompt_tokens: u32,
pub completion_tokens: u32,
pub total_tokens: u32,
}
// -- Handlers --
/// Phase 39: resolve (provider, effective_model) from a ChatRequest.
///
/// Explicit `req.provider` wins. If absent, infer from a model-name
/// prefix: "openrouter/..." → openrouter (strip prefix), "cloud/..." →
/// ollama_cloud (strip prefix). Bare names default to "ollama".
///
/// The stripped model is what the upstream adapter expects:
/// OpenRouter's API wants "openai/gpt-4o-mini", not
/// "openrouter/openai/gpt-4o-mini".
fn resolve_provider(req: &ChatRequest) -> (String, String) {
if let Some(p) = req.provider.as_deref() {
return (p.to_ascii_lowercase(), req.model.clone());
}
if let Some(rest) = req.model.strip_prefix("openrouter/") {
return ("openrouter".to_string(), rest.to_string());
}
if let Some(rest) = req.model.strip_prefix("cloud/") {
return ("ollama_cloud".to_string(), rest.to_string());
}
if let Some(rest) = req.model.strip_prefix("gemini/") {
return ("gemini".to_string(), rest.to_string());
}
if let Some(rest) = req.model.strip_prefix("claude/") {
return ("claude".to_string(), rest.to_string());
}
("ollama".to_string(), req.model.clone())
}
#[cfg(test)]
mod resolve_provider_tests {
use super::*;
fn mk_req(provider: Option<&str>, model: &str) -> ChatRequest {
ChatRequest {
model: model.to_string(),
messages: vec![],
temperature: None,
max_tokens: None,
stream: None,
think: None,
provider: provider.map(|s| s.to_string()),
}
}
#[test]
fn explicit_provider_wins() {
let r = mk_req(Some("openrouter"), "qwen3.5:latest");
assert_eq!(resolve_provider(&r), ("openrouter".into(), "qwen3.5:latest".into()));
}
#[test]
fn bare_model_defaults_to_ollama() {
let r = mk_req(None, "qwen3.5:latest");
assert_eq!(resolve_provider(&r), ("ollama".into(), "qwen3.5:latest".into()));
}
#[test]
fn openrouter_prefix_infers_and_strips() {
let r = mk_req(None, "openrouter/openai/gpt-4o-mini");
assert_eq!(resolve_provider(&r), ("openrouter".into(), "openai/gpt-4o-mini".into()));
}
#[test]
fn cloud_prefix_infers_and_strips() {
let r = mk_req(None, "cloud/kimi-k2:1t");
assert_eq!(resolve_provider(&r), ("ollama_cloud".into(), "kimi-k2:1t".into()));
}
#[test]
fn explicit_provider_preserves_full_model_even_with_prefix() {
// If caller provides both provider and a model with a prefix,
// trust them — don't strip. The adapter will get the full model
// string as-is.
let r = mk_req(Some("openrouter"), "openrouter/openai/gpt-4o-mini");
assert_eq!(resolve_provider(&r), ("openrouter".into(), "openrouter/openai/gpt-4o-mini".into()));
}
#[test]
fn gemini_prefix_infers_and_strips() {
let r = mk_req(None, "gemini/gemini-2.0-flash");
assert_eq!(resolve_provider(&r), ("gemini".into(), "gemini-2.0-flash".into()));
}
#[test]
fn claude_prefix_infers_and_strips() {
let r = mk_req(None, "claude/claude-3-5-sonnet-latest");
assert_eq!(resolve_provider(&r), ("claude".into(), "claude-3-5-sonnet-latest".into()));
}
}
async fn chat(
State(state): State<V1State>,
Json(req): Json<ChatRequest>,
) -> Result<Json<ChatResponse>, (StatusCode, String)> {
if req.messages.is_empty() {
return Err((StatusCode::BAD_REQUEST, "messages must be non-empty".into()));
}
if req.stream.unwrap_or(false) {
tracing::warn!("/v1/chat: stream=true requested but Phase 38 returns non-streaming");
}
// Provider resolution: explicit `req.provider` wins; otherwise
// infer from a model-name prefix. Phase 39 PRD gate example:
// `model: "openrouter/openai/gpt-4o-mini"` → provider "openrouter",
// adapter gets the stripped "openai/gpt-4o-mini".
let (provider, effective_model) = resolve_provider(&req);
let start_time = chrono::Utc::now();
let start_instant = std::time::Instant::now();
// If we stripped a prefix, clone req with the effective model so
// the adapter sees what the upstream provider expects (OpenRouter
// wants "openai/gpt-4o-mini", not "openrouter/openai/gpt-4o-mini").
let req_for_adapter: std::borrow::Cow<'_, ChatRequest> =
if effective_model == req.model {
std::borrow::Cow::Borrowed(&req)
} else {
let mut cloned = req.clone();
cloned.model = effective_model.clone();
std::borrow::Cow::Owned(cloned)
};
let (resp, used_provider) = match provider.as_str() {
"ollama" | "local" | "" => {
let r = ollama::chat(&state.ai_client, &*req_for_adapter)
.await
.map_err(|e| (StatusCode::BAD_GATEWAY, format!("ollama local: {e}")))?;
(r, "ollama".to_string())
}
"ollama_cloud" | "cloud" => {
let key = state.ollama_cloud_key.as_deref().ok_or((
StatusCode::SERVICE_UNAVAILABLE,
"OLLAMA_CLOUD_KEY not configured".to_string(),
))?;
let r = ollama_cloud::chat(key, &*req_for_adapter)
.await
.map_err(|e| (StatusCode::BAD_GATEWAY, format!("ollama cloud: {e}")))?;
(r, "ollama_cloud".to_string())
}
"openrouter" | "openrouter_free" => {
// Free-tier rescue rung. Added 2026-04-24 after iter 5
// repeated Ollama Cloud 502s on kimi-k2:1t — OpenRouter
// gives a different provider backbone as fallback.
let key = state.openrouter_key.as_deref().ok_or((
StatusCode::SERVICE_UNAVAILABLE,
"OPENROUTER_API_KEY not configured".to_string(),
))?;
let r = openrouter::chat(key, &*req_for_adapter)
.await
.map_err(|e| (StatusCode::BAD_GATEWAY, format!("openrouter: {e}")))?;
(r, "openrouter".to_string())
}
"gemini" => {
// Phase 40 provider adapter. Google Generative Language
// API via query-string key auth (not bearer).
let key = state.gemini_key.as_deref().ok_or((
StatusCode::SERVICE_UNAVAILABLE,
"GEMINI_API_KEY not configured".to_string(),
))?;
let r = gemini::chat(key, &*req_for_adapter)
.await
.map_err(|e| (StatusCode::BAD_GATEWAY, format!("gemini: {e}")))?;
(r, "gemini".to_string())
}
"claude" | "anthropic" => {
// Phase 40 provider adapter. Anthropic Messages API via
// x-api-key header + anthropic-version:2023-06-01.
let key = state.claude_key.as_deref().ok_or((
StatusCode::SERVICE_UNAVAILABLE,
"ANTHROPIC_API_KEY not configured".to_string(),
))?;
let r = claude::chat(key, &*req_for_adapter)
.await
.map_err(|e| (StatusCode::BAD_GATEWAY, format!("claude: {e}")))?;
(r, "claude".to_string())
}
other => {
return Err((
StatusCode::BAD_REQUEST,
format!("unknown provider '{other}' — supported: ollama, ollama_cloud, openrouter, gemini, claude"),
));
}
};
let end_time = chrono::Utc::now();
let latency_ms = start_instant.elapsed().as_millis() as u64;
// Phase 40 — emit Langfuse trace. Fire-and-forget: the clone is
// cheap (Arc inside), the tokio::spawn never blocks us, a dead
// Langfuse just logs a warn. Client-visible response latency is
// untouched.
if let Some(lf) = &state.langfuse {
let output = resp.choices.first()
.map(|c| c.message.content.clone())
.unwrap_or_default();
lf.emit_chat(langfuse_trace::ChatTrace {
provider: used_provider.clone(),
model: resp.model.clone(),
input: req.messages.clone(),
output,
prompt_tokens: resp.usage.prompt_tokens,
completion_tokens: resp.usage.completion_tokens,
temperature: req.temperature,
max_tokens: req.max_tokens,
think: req.think,
start_time: start_time.to_rfc3339(),
end_time: end_time.to_rfc3339(),
latency_ms,
});
}
// Phase 40: per-provider usage tracking
{
let mut u = state.usage.write().await;
u.requests += 1;
u.prompt_tokens += resp.usage.prompt_tokens as u64;
u.completion_tokens += resp.usage.completion_tokens as u64;
u.total_tokens += resp.usage.total_tokens as u64;
let provider_usage = u.by_provider.entry(used_provider).or_default();
provider_usage.requests += 1;
provider_usage.prompt_tokens += resp.usage.prompt_tokens as u64;
provider_usage.completion_tokens += resp.usage.completion_tokens as u64;
provider_usage.total_tokens += resp.usage.total_tokens as u64;
}
Ok(Json(resp))
}
async fn usage(State(state): State<V1State>) -> impl IntoResponse {
let snapshot = state.usage.read().await.clone();
Json(snapshot)
}
// Phase 38 is stateless — no session persistence yet. Return an empty
// list in OpenAI-ish shape so clients that probe this endpoint don't
// 404. Real session state lands in Phase 41 with the profile-system
// expansion.
async fn sessions() -> impl IntoResponse {
Json(serde_json::json!({
"data": [],
"object": "list",
"note": "Phase 38: stateless. Session state lands in Phase 41.",
}))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn chat_request_parses_openai_shape() {
let raw = r#"{
"model": "qwen3.5:latest",
"messages": [
{"role": "system", "content": "You are helpful."},
{"role": "user", "content": "Hi"}
],
"temperature": 0.2,
"max_tokens": 100
}"#;
let r: ChatRequest = serde_json::from_str(raw).unwrap();
assert_eq!(r.model, "qwen3.5:latest");
assert_eq!(r.messages.len(), 2);
assert_eq!(r.messages[0].role, "system");
assert_eq!(r.messages[1].content, "Hi");
assert_eq!(r.temperature, Some(0.2));
assert_eq!(r.max_tokens, Some(100));
}
#[test]
fn chat_request_accepts_minimal() {
let raw = r#"{
"model": "any",
"messages": [{"role": "user", "content": "hi"}]
}"#;
let r: ChatRequest = serde_json::from_str(raw).unwrap();
assert_eq!(r.temperature, None);
assert_eq!(r.max_tokens, None);
assert_eq!(r.stream, None);
}
#[test]
fn usage_counter_default_is_zero() {
let u = Usage::default();
assert_eq!(u.requests, 0);
assert_eq!(u.total_tokens, 0);
}
}
+160
View File
@@ -0,0 +1,160 @@
//! Phase 38 — Ollama shape adapter.
//!
//! Translates `/v1/chat` (OpenAI-compatible) requests into the
//! existing aibridge `GenerateRequest` shape, and the `GenerateResponse`
//! back into an OpenAI-compatible `ChatResponse`. This is a bridge,
//! not a new client — aibridge + the Python sidecar stay as-is.
//!
//! Phase 39 replaces this direct call with a `ProviderAdapter` trait
//! dispatch so the same `/v1/chat` handler routes to any provider.
use aibridge::client::{AiClient, GenerateRequest};
use super::{ChatRequest, ChatResponse, Choice, Message, UsageBlock};
pub async fn chat(client: &AiClient, req: &ChatRequest) -> Result<ChatResponse, String> {
let (system, prompt) = flatten_messages(&req.messages);
let gen_req = GenerateRequest {
prompt,
model: Some(req.model.clone()),
system: if system.is_empty() { None } else { Some(system) },
temperature: req.temperature,
max_tokens: req.max_tokens,
// Phase 38 default: think=false. Hot-path discipline for
// thinking models (qwen3.5, qwen3, gpt-oss) which otherwise
// burn the max_tokens budget on hidden reasoning before any
// visible output, producing empty responses. Callers that
// actually want reasoning (overseers, T3+ tiers) opt in via
// the `think: true` extension field. Phase 40 routing engine
// flips this per task class.
think: Some(req.think.unwrap_or(false)),
};
let t0 = std::time::Instant::now();
let resp = client.generate(gen_req).await.map_err(|e| e.to_string())?;
let latency_ms = t0.elapsed().as_millis();
// Prefer sidecar-reported token counts when present. Fall back to
// chars/4 estimate (biased safe ~15%, matches Phase 21 convention
// in crates/aibridge/src/context.rs::estimate_tokens).
let prompt_tokens = resp.tokens_evaluated
.map(|n| n as u32)
.unwrap_or_else(|| estimate_prompt_tokens(&req.messages));
let completion_tokens = resp.tokens_generated
.map(|n| n as u32)
.unwrap_or_else(|| ((resp.text.chars().count() + 3) / 4) as u32);
tracing::info!(
target: "v1.chat",
model = %req.model,
prompt_tokens,
completion_tokens,
latency_ms = latency_ms as u64,
"ollama chat completed",
);
Ok(ChatResponse {
id: format!("chatcmpl-{}", chrono::Utc::now().timestamp_nanos_opt().unwrap_or(0)),
object: "chat.completion",
created: chrono::Utc::now().timestamp(),
model: resp.model,
choices: vec![Choice {
index: 0,
message: Message {
role: "assistant".into(),
content: resp.text,
},
finish_reason: "stop".into(),
}],
usage: UsageBlock {
prompt_tokens,
completion_tokens,
total_tokens: prompt_tokens + completion_tokens,
},
})
}
/// Public re-export of the flattener so sibling adapters (Ollama Cloud,
/// future OpenRouter) can reuse the same shape collapse without
/// duplicating the logic. Keeps `(system, prompt)` format consistent
/// across providers.
pub fn flatten_messages_public(messages: &[Message]) -> (String, String) {
flatten_messages(messages)
}
/// Collapse a message array into (system, prompt). Multiple system
/// messages concatenate with a newline — matches OpenAI's documented
/// behavior. Non-system messages become role-labeled turns.
fn flatten_messages(messages: &[Message]) -> (String, String) {
let mut system = String::new();
let mut prompt = String::new();
for m in messages {
if m.role == "system" {
if !system.is_empty() { system.push('\n'); }
system.push_str(&m.content);
} else {
prompt.push_str(&m.role);
prompt.push_str(": ");
prompt.push_str(&m.content);
prompt.push_str("\n\n");
}
}
prompt.push_str("assistant:");
(system, prompt)
}
fn estimate_prompt_tokens(messages: &[Message]) -> u32 {
let chars: usize = messages.iter().map(|m| m.content.chars().count()).sum();
((chars + 3) / 4) as u32
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn flatten_separates_system_from_turns() {
let msgs = vec![
Message { role: "system".into(), content: "Rules here.".into() },
Message { role: "user".into(), content: "Q1".into() },
Message { role: "assistant".into(), content: "A1".into() },
Message { role: "user".into(), content: "Q2".into() },
];
let (system, prompt) = flatten_messages(&msgs);
assert_eq!(system, "Rules here.");
assert!(prompt.contains("user: Q1"));
assert!(prompt.contains("assistant: A1"));
assert!(prompt.contains("user: Q2"));
assert!(prompt.trim_end().ends_with("assistant:"));
}
#[test]
fn flatten_concatenates_multiple_system_messages() {
let msgs = vec![
Message { role: "system".into(), content: "First.".into() },
Message { role: "system".into(), content: "Second.".into() },
Message { role: "user".into(), content: "Hi".into() },
];
let (system, _) = flatten_messages(&msgs);
assert_eq!(system, "First.\nSecond.");
}
#[test]
fn flatten_with_no_system_returns_empty_system() {
let msgs = vec![Message { role: "user".into(), content: "hi".into() }];
let (system, prompt) = flatten_messages(&msgs);
assert!(system.is_empty());
assert!(prompt.contains("user: hi"));
}
#[test]
fn estimate_tokens_chars_div_4_ceiling() {
let msgs = vec![Message { role: "user".into(), content: "abcdefgh".into() }];
// 8 chars / 4 = 2, with ceiling → 2
assert_eq!(estimate_prompt_tokens(&msgs), 2);
let msgs2 = vec![Message { role: "user".into(), content: "abcdefghi".into() }];
// 9 chars → (9+3)/4 = 3 (ceiling)
assert_eq!(estimate_prompt_tokens(&msgs2), 3);
}
}
+189
View File
@@ -0,0 +1,189 @@
//! Phase 39 (first slice) — Ollama Cloud adapter.
//!
//! Direct HTTPS call to `https://ollama.com/api/generate` with Bearer
//! auth. Mirrors the `tests/multi-agent/agent.ts::generateCloud` pattern
//! exactly (same endpoint, same body shape, same Bearer header), so
//! cloud calls from Rust behave identically to the TS hot path.
//!
//! Key sourcing priority:
//! 1. Env var `OLLAMA_CLOUD_KEY` (matches agent.ts convention)
//! 2. `/root/llm_team_config.json` → providers.ollama_cloud.api_key
//! 3. Env var `OLLAMA_CLOUD_API_KEY` (LLM Team UI convention)
//!
//! First hit wins. Key is loaded once at gateway startup by
//! `resolve_cloud_key()` and stored on `V1State.ollama_cloud_key`.
use std::time::Duration;
use serde::{Deserialize, Serialize};
use super::{ChatRequest, ChatResponse, Choice, Message, UsageBlock};
const CLOUD_BASE_URL: &str = "https://ollama.com";
const CLOUD_TIMEOUT_SECS: u64 = 180;
/// Read the Ollama Cloud key from the three sanctioned sources. Returns
/// None if none is set — callers must 503 rather than attempt a call.
pub fn resolve_cloud_key() -> Option<String> {
if let Ok(k) = std::env::var("OLLAMA_CLOUD_KEY") {
if !k.trim().is_empty() { return Some(k.trim().to_string()); }
}
if let Ok(raw) = std::fs::read_to_string("/root/llm_team_config.json") {
if let Ok(v) = serde_json::from_str::<serde_json::Value>(&raw) {
if let Some(k) = v.pointer("/providers/ollama_cloud/api_key").and_then(|x| x.as_str()) {
if !k.trim().is_empty() { return Some(k.trim().to_string()); }
}
}
}
if let Ok(k) = std::env::var("OLLAMA_CLOUD_API_KEY") {
if !k.trim().is_empty() { return Some(k.trim().to_string()); }
}
None
}
pub async fn chat(
key: &str,
req: &ChatRequest,
) -> Result<ChatResponse, String> {
let (system, prompt) = super::ollama::flatten_messages_public(&req.messages);
let body = CloudGenerateBody {
model: req.model.clone(),
prompt,
system: if system.is_empty() { None } else { Some(system) },
stream: false,
think: Some(req.think.unwrap_or(false)),
options: CloudOptions {
// Thinking cloud models need headroom — floor 400 to give
// qwen3.5:397b / gpt-oss:120b reasoning room. Matches
// agent.ts `Math.max(opts.max_tokens ?? 800, 400)` policy.
num_predict: std::cmp::max(req.max_tokens.unwrap_or(800), 400),
temperature: req.temperature.unwrap_or(0.3),
},
};
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(CLOUD_TIMEOUT_SECS))
.build()
.map_err(|e| format!("build client: {e}"))?;
let t0 = std::time::Instant::now();
let resp = client
.post(format!("{}/api/generate", CLOUD_BASE_URL))
.bearer_auth(key)
.json(&body)
.send()
.await
.map_err(|e| format!("ollama.com unreachable: {e}"))?;
let status = resp.status();
if !status.is_success() {
let body = resp.text().await.unwrap_or_else(|_| "?".into());
return Err(format!("ollama.com {}: {}", status, body));
}
let parsed: CloudGenerateResponse = resp.json().await
.map_err(|e| format!("invalid cloud response: {e}"))?;
let latency_ms = t0.elapsed().as_millis();
let text = parsed.response.unwrap_or_default();
let prompt_tokens = parsed.prompt_eval_count.unwrap_or_else(|| {
let chars: usize = req.messages.iter().map(|m| m.content.chars().count()).sum();
((chars + 3) / 4) as u32
});
let completion_tokens = parsed.eval_count.unwrap_or_else(|| {
((text.chars().count() + 3) / 4) as u32
});
tracing::info!(
target: "v1.chat",
provider = "ollama_cloud",
model = %req.model,
prompt_tokens,
completion_tokens,
latency_ms = latency_ms as u64,
"cloud chat completed",
);
Ok(ChatResponse {
id: format!("chatcmpl-{}", chrono::Utc::now().timestamp_nanos_opt().unwrap_or(0)),
object: "chat.completion",
created: chrono::Utc::now().timestamp(),
model: parsed.model.unwrap_or_else(|| req.model.clone()),
choices: vec![Choice {
index: 0,
message: Message { role: "assistant".into(), content: text },
finish_reason: "stop".into(),
}],
usage: UsageBlock {
prompt_tokens,
completion_tokens,
total_tokens: prompt_tokens + completion_tokens,
},
})
}
// -- Ollama /api/generate wire shapes --
#[derive(Serialize)]
struct CloudGenerateBody {
model: String,
prompt: String,
#[serde(skip_serializing_if = "Option::is_none")]
system: Option<String>,
stream: bool,
#[serde(skip_serializing_if = "Option::is_none")]
think: Option<bool>,
options: CloudOptions,
}
#[derive(Serialize)]
struct CloudOptions {
num_predict: u32,
temperature: f64,
}
#[derive(Deserialize)]
struct CloudGenerateResponse {
#[serde(default)]
response: Option<String>,
#[serde(default)]
model: Option<String>,
#[serde(default)]
prompt_eval_count: Option<u32>,
#[serde(default)]
eval_count: Option<u32>,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn resolve_cloud_key_returns_none_when_no_sources_set() {
// Only check shape — we can't reliably unset env vars in a test
// that runs alongside others, and the file path is on disk.
// If all three sources are empty the function returns None; if
// any are set we expect Some. This just smoke-tests the call
// doesn't panic.
let _ = resolve_cloud_key();
}
#[test]
fn cloud_body_serializes_compact() {
let body = CloudGenerateBody {
model: "gpt-oss:120b".into(),
prompt: "user: hi\n\nassistant:".into(),
system: Some("Be terse.".into()),
stream: false,
think: Some(false),
options: CloudOptions { num_predict: 400, temperature: 0.3 },
};
let json = serde_json::to_string(&body).unwrap();
assert!(json.contains("\"model\":\"gpt-oss:120b\""));
assert!(json.contains("\"stream\":false"));
assert!(json.contains("\"num_predict\":400"));
assert!(json.contains("\"think\":false"));
assert!(json.contains("\"system\":\"Be terse.\""));
}
}
+217
View File
@@ -0,0 +1,217 @@
//! OpenRouter adapter — free-tier rescue rung for /v1/chat.
//!
//! Direct HTTPS call to `https://openrouter.ai/api/v1/chat/completions`
//! with Bearer auth. Mirrors the OpenAI-compatible shape so the model
//! list can be expanded without code changes. Added 2026-04-24 after
//! iter 5 hit repeated Ollama Cloud 502s on kimi-k2:1t — OpenRouter
//! free-tier models give us a different provider backbone as fallback.
//!
//! Key sourcing priority:
//! 1. Env var `OPENROUTER_API_KEY`
//! 2. `/home/profit/.env` (LLM Team convention)
//! 3. `/root/llm_team_config.json` → providers.openrouter.api_key
//!
//! First hit wins. Key is resolved once at gateway startup and stored
//! on `V1State.openrouter_key`.
use std::time::Duration;
use serde::{Deserialize, Serialize};
use super::{ChatRequest, ChatResponse, Choice, Message, UsageBlock};
const OR_BASE_URL: &str = "https://openrouter.ai/api/v1";
const OR_TIMEOUT_SECS: u64 = 180;
pub fn resolve_openrouter_key() -> Option<String> {
if let Ok(k) = std::env::var("OPENROUTER_API_KEY") {
if !k.trim().is_empty() { return Some(k.trim().to_string()); }
}
// LLM Team UI writes its key to ~/.env on the host user — pick it up
// from the same source so the free-tier rescue path works without
// an explicit systemd Environment= line.
for path in ["/home/profit/.env", "/root/.env"] {
if let Ok(raw) = std::fs::read_to_string(path) {
for line in raw.lines() {
if let Some(rest) = line.strip_prefix("OPENROUTER_API_KEY=") {
let k = rest.trim().trim_matches('"').trim_matches('\'');
if !k.is_empty() { return Some(k.to_string()); }
}
}
}
}
if let Ok(raw) = std::fs::read_to_string("/root/llm_team_config.json") {
if let Ok(v) = serde_json::from_str::<serde_json::Value>(&raw) {
if let Some(k) = v.pointer("/providers/openrouter/api_key").and_then(|x| x.as_str()) {
if !k.trim().is_empty() { return Some(k.trim().to_string()); }
}
}
}
None
}
pub async fn chat(
key: &str,
req: &ChatRequest,
) -> Result<ChatResponse, String> {
// Strip the "openrouter/" prefix if the caller used the namespaced
// form so OpenRouter sees the raw model id (e.g. "openai/gpt-oss-120b:free").
let model = req.model.strip_prefix("openrouter/").unwrap_or(&req.model).to_string();
let body = ORChatBody {
model: model.clone(),
messages: req.messages.iter().map(|m| ORMessage {
role: m.role.clone(),
content: m.content.clone(),
}).collect(),
max_tokens: req.max_tokens.unwrap_or(800),
temperature: req.temperature.unwrap_or(0.3),
stream: false,
};
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(OR_TIMEOUT_SECS))
.build()
.map_err(|e| format!("build client: {e}"))?;
let t0 = std::time::Instant::now();
let resp = client
.post(format!("{}/chat/completions", OR_BASE_URL))
.bearer_auth(key)
// OpenRouter recommends Referer + Title for attribution; absent
// headers do not fail the call but help us see our traffic in
// their dashboard.
.header("HTTP-Referer", "https://vcp.devop.live")
.header("X-Title", "Lakehouse Scrum")
.json(&body)
.send()
.await
.map_err(|e| format!("openrouter.ai unreachable: {e}"))?;
let status = resp.status();
if !status.is_success() {
let body = resp.text().await.unwrap_or_else(|_| "?".into());
return Err(format!("openrouter.ai {}: {}", status, body));
}
let parsed: ORChatResponse = resp.json().await
.map_err(|e| format!("invalid openrouter response: {e}"))?;
let latency_ms = t0.elapsed().as_millis();
let choice = parsed.choices.into_iter().next()
.ok_or_else(|| "openrouter returned no choices".to_string())?;
let text = choice.message.content;
let prompt_tokens = parsed.usage.as_ref().map(|u| u.prompt_tokens).unwrap_or_else(|| {
let chars: usize = req.messages.iter().map(|m| m.content.chars().count()).sum();
((chars + 3) / 4) as u32
});
let completion_tokens = parsed.usage.as_ref().map(|u| u.completion_tokens).unwrap_or_else(|| {
((text.chars().count() + 3) / 4) as u32
});
tracing::info!(
target: "v1.chat",
provider = "openrouter",
model = %model,
prompt_tokens,
completion_tokens,
latency_ms = latency_ms as u64,
"openrouter chat completed",
);
Ok(ChatResponse {
id: format!("chatcmpl-{}", chrono::Utc::now().timestamp_nanos_opt().unwrap_or(0)),
object: "chat.completion",
created: chrono::Utc::now().timestamp(),
model,
choices: vec![Choice {
index: 0,
message: Message { role: "assistant".into(), content: text },
finish_reason: choice.finish_reason.unwrap_or_else(|| "stop".into()),
}],
usage: UsageBlock {
prompt_tokens,
completion_tokens,
total_tokens: prompt_tokens + completion_tokens,
},
})
}
// -- OpenRouter wire shapes (OpenAI-compatible) --
#[derive(Serialize)]
struct ORChatBody {
model: String,
messages: Vec<ORMessage>,
max_tokens: u32,
temperature: f64,
stream: bool,
}
#[derive(Serialize)]
struct ORMessage { role: String, content: String }
#[derive(Deserialize)]
struct ORChatResponse {
choices: Vec<ORChoice>,
#[serde(default)]
usage: Option<ORUsage>,
}
#[derive(Deserialize)]
struct ORChoice {
message: ORMessageResp,
#[serde(default)]
finish_reason: Option<String>,
}
#[derive(Deserialize)]
struct ORMessageResp { content: String }
#[derive(Deserialize)]
struct ORUsage { prompt_tokens: u32, completion_tokens: u32 }
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn resolve_openrouter_key_does_not_panic() {
// Smoke test — all three sources may or may not be set depending
// on environment; just confirm the call returns cleanly.
let _ = resolve_openrouter_key();
}
#[test]
fn chat_body_serializes_to_openai_shape() {
let body = ORChatBody {
model: "openai/gpt-oss-120b:free".into(),
messages: vec![
ORMessage { role: "user".into(), content: "review this".into() },
],
max_tokens: 800,
temperature: 0.3,
stream: false,
};
let json = serde_json::to_string(&body).unwrap();
assert!(json.contains("\"model\":\"openai/gpt-oss-120b:free\""));
assert!(json.contains("\"messages\""));
assert!(json.contains("\"max_tokens\":800"));
assert!(json.contains("\"stream\":false"));
}
#[test]
fn model_prefix_strip_preserves_unprefixed() {
// If caller passes "openrouter/openai/gpt-oss-120b:free" we strip.
// If caller passes "openai/gpt-oss-120b:free" unchanged, we keep.
let cases = [
("openrouter/openai/gpt-oss-120b:free", "openai/gpt-oss-120b:free"),
("openai/gpt-oss-120b:free", "openai/gpt-oss-120b:free"),
("google/gemma-3-27b-it:free", "google/gemma-3-27b-it:free"),
];
for (input, expected) in cases {
let out = input.strip_prefix("openrouter/").unwrap_or(input);
assert_eq!(out, expected, "{input} should become {expected}");
}
}
}
+150
View File
@@ -0,0 +1,150 @@
//! `/v1/respond` — the **execution** API (distinct from `/v1/chat`, the
//! completion API).
//!
//! This is the consolidation move called out in the 2026-04-23 session:
//! lift the proven pipeline from `tests/multi-agent/orchestrator.ts`
//! (executor → reviewer → escalate → validate → seal playbook →
//! write-through to KB) into the gateway, so the production path has
//! the intelligence the tests already proved.
//!
//! `/v1/chat` stays a naive completion proxy for callers that want one.
//! `/v1/respond` is where the loop lives. Every orchestrator-style
//! caller migrates here and the TS harnesses become thin clients.
//!
//! This file holds the HTTP surface + request/response shapes. The loop
//! itself lives in `execution_loop::ExecutionLoop`.
use axum::{extract::State, http::StatusCode, Json};
use serde::{Deserialize, Serialize};
use super::V1State;
use crate::execution_loop::{ExecutionLoop, LogEntry, RespondOutcome};
/// A structured task — mirrors `TaskSpec` in `tests/multi-agent/agent.ts`.
/// Kept deliberately open so non-staffing task classes (code-gen,
/// DevOps-long-horizon) can land without a schema fight.
#[derive(Deserialize, Debug, Clone)]
pub struct RespondRequest {
/// Task class — routes to the right truth rules + validator. For the
/// staffing substrate: `staffing.fill`, `staffing.rescue`,
/// `staffing.sms_draft`. Truth-layer lookup is a no-op until a rule
/// set is registered for the class.
pub task_class: String,
/// Human-readable operation description — becomes the playbook
/// `operation` field on seal, and the primary signal for
/// playbook_memory embedding.
pub operation: String,
/// Free-form structured context. Passed to the executor prompt and
/// to the playbook seeder. Staffing tasks expect
/// `{target_role, target_count, target_city, target_state, approach_hint}`
/// but nothing here validates that — the validator crate will (Phase 43).
#[serde(default)]
pub spec: serde_json::Value,
/// Executor model. Defaults to the hot-path local model if omitted.
/// See orchestrator.ts:28 (`EXECUTOR_MODEL = "qwen3.5:latest"`).
#[serde(default)]
pub executor_model: Option<String>,
/// Reviewer model. Defaults to the hot-path local reviewer.
/// See orchestrator.ts:29 (`REVIEWER_MODEL = "qwen3:latest"`).
#[serde(default)]
pub reviewer_model: Option<String>,
/// Hard cap on executor turns. Default matches orchestrator.ts:30
/// (`MAX_TURNS = 12`). Cloud escalation counts as a turn.
#[serde(default)]
pub max_turns: Option<u32>,
}
#[derive(Serialize)]
pub struct RespondResponse {
/// `ok` = consensus reached, playbook sealed. `failed` = loop ran
/// out of turns or hit the drift cap. `blocked` = truth-layer
/// veto (Phase 42 rule citation in `error`).
pub status: &'static str,
/// The final artifact — for staffing fills, `{fills: [{candidate_id, name}]}`.
/// Empty on failure / block.
pub artifact: serde_json::Value,
/// Structured cross-turn log. Same shape as orchestrator.ts LogEntry
/// so existing tooling (kb extractors, fact_extractor.ts) reads it
/// without change.
pub log: Vec<LogEntry>,
/// Iteration count actually used. ≤ max_turns. Stamped on
/// outcomes.jsonl per the indicator audit (2026-04-23).
pub iterations: u32,
/// Error message on non-ok status. Truth-rule citations land here
/// when `status == "blocked"`.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
}
pub async fn respond(
State(state): State<V1State>,
Json(req): Json<RespondRequest>,
) -> Result<Json<RespondResponse>, (StatusCode, String)> {
if req.operation.is_empty() {
return Err((StatusCode::BAD_REQUEST, "operation must be non-empty".into()));
}
if req.task_class.is_empty() {
return Err((StatusCode::BAD_REQUEST, "task_class must be non-empty".into()));
}
let mut loop_runner = ExecutionLoop::new(state, req);
let outcome = loop_runner.run().await.map_err(|e| {
(StatusCode::INTERNAL_SERVER_ERROR, format!("execution loop: {e}"))
})?;
let (status, error) = match &outcome {
RespondOutcome::Ok { .. } => ("ok", None),
RespondOutcome::Failed { reason, .. } => ("failed", Some(reason.clone())),
RespondOutcome::Blocked { reason, .. } => ("blocked", Some(reason.clone())),
};
Ok(Json(RespondResponse {
status,
artifact: outcome.artifact(),
log: outcome.into_log(),
iterations: loop_runner.turns_used(),
error,
}))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn respond_request_parses_minimal() {
let raw = r#"{
"task_class": "staffing.fill",
"operation": "fill: Welder x2 in Toledo, OH"
}"#;
let r: RespondRequest = serde_json::from_str(raw).unwrap();
assert_eq!(r.task_class, "staffing.fill");
assert_eq!(r.executor_model, None);
assert_eq!(r.max_turns, None);
}
#[test]
fn respond_request_parses_full() {
let raw = r#"{
"task_class": "staffing.fill",
"operation": "fill: Welder x2 in Toledo, OH",
"spec": {"target_role": "Welder", "target_count": 2, "target_city": "Toledo", "target_state": "OH"},
"executor_model": "qwen3.5:latest",
"reviewer_model": "qwen3:latest",
"max_turns": 12
}"#;
let r: RespondRequest = serde_json::from_str(raw).unwrap();
assert_eq!(r.executor_model.as_deref(), Some("qwen3.5:latest"));
assert_eq!(r.max_turns, Some(12));
assert_eq!(r.spec["target_count"], 2);
}
}
+47
View File
@@ -0,0 +1,47 @@
use serde::Serialize;
use truth::default_truth_store;
// Note: truth_router() was a stub wrapper around a single /context route
// that nothing called — v1/mod.rs wires get(truth::context) directly
// onto its own router. Removed 2026-04-24 along with its #[allow(dead_code)]
// attribute; the handler below is the real surface.
#[derive(Serialize)]
pub struct ContextResponse {
pub task_classes: Vec<String>,
pub rules: Vec<RuleInfo>,
}
#[derive(Serialize)]
pub struct RuleInfo {
pub id: String,
pub task_class: String,
pub description: String,
}
pub async fn context() -> axum::Json<ContextResponse> {
let store = default_truth_store();
let task_classes: Vec<String> = vec![
"staffing.fill".to_string(),
"staffing.rescue".to_string(),
"staffing.sms_draft".to_string(),
"staffing.any".to_string(),
];
let mut rules = Vec::new();
for tc in &task_classes {
for rule in store.get_rules(tc) {
rules.push(RuleInfo {
id: rule.id.clone(),
task_class: rule.task_class.clone(),
description: rule.description.clone(),
});
}
}
axum::Json(ContextResponse {
task_classes,
rules,
})
}
+28
View File
@@ -0,0 +1,28 @@
[package]
name = "ingestd"
version = "0.1.0"
edition = "2024"
[dependencies]
shared = { path = "../shared" }
storaged = { path = "../storaged" }
catalogd = { path = "../catalogd" }
vectord = { path = "../vectord" }
journald = { path = "../journald" }
tokio = { workspace = true }
axum = { workspace = true, features = ["multipart"] }
lopdf = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }
tracing = { workspace = true }
arrow = { workspace = true }
parquet = { workspace = true }
bytes = { workspace = true }
sha2 = { workspace = true }
csv = { workspace = true }
chrono = { workspace = true }
object_store = { workspace = true }
tokio-postgres = { workspace = true }
mysql_async = { workspace = true }
uuid = { workspace = true }
croner = "3"
+203
View File
@@ -0,0 +1,203 @@
use arrow::array::{ArrayRef, Float64Array, Int64Array, StringArray, BooleanArray};
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use std::sync::Arc;
/// Inferred column type from sampling data.
#[derive(Debug, Clone, PartialEq)]
enum InferredType {
Integer,
Float,
Boolean,
String,
}
/// Parse CSV bytes into Arrow RecordBatches with automatic schema detection.
/// Per ADR-010: ambiguous types default to String.
pub fn parse_csv(content: &[u8]) -> Result<(Arc<Schema>, Vec<RecordBatch>), String> {
let mut reader = csv::ReaderBuilder::new()
.flexible(true) // allow varying column counts
.trim(csv::Trim::All)
.from_reader(content);
let headers: Vec<String> = reader.headers()
.map_err(|e| format!("CSV header error: {e}"))?
.iter()
.enumerate()
.map(|(i, h)| {
let h = h.trim().to_string();
if h.is_empty() { format!("column_{i}") } else { sanitize_column_name(&h) }
})
.collect();
let n_cols = headers.len();
if n_cols == 0 {
return Err("CSV has no columns".into());
}
// Read all rows into string columns
let mut columns: Vec<Vec<String>> = vec![vec![]; n_cols];
let mut row_count = 0;
for result in reader.records() {
let record = result.map_err(|e| format!("CSV row error: {e}"))?;
for (i, field) in record.iter().enumerate() {
if i < n_cols {
columns[i].push(field.trim().to_string());
}
}
// Pad short rows with empty strings
for col in columns.iter_mut().skip(record.len().min(n_cols)) {
col.push(String::new());
}
row_count += 1;
}
if row_count == 0 {
return Err("CSV has no data rows".into());
}
tracing::info!("parsed CSV: {row_count} rows × {n_cols} columns");
// Infer types by sampling (look at all values)
let types: Vec<InferredType> = columns.iter().map(|col| infer_column_type(col)).collect();
// Build Arrow schema
let fields: Vec<Field> = headers.iter().zip(types.iter()).map(|(name, typ)| {
let dt = match typ {
InferredType::Integer => DataType::Int64,
InferredType::Float => DataType::Float64,
InferredType::Boolean => DataType::Boolean,
InferredType::String => DataType::Utf8,
};
Field::new(name, dt, true) // all nullable
}).collect();
let schema = Arc::new(Schema::new(fields));
// Build arrays
let arrays: Vec<ArrayRef> = columns.iter().zip(types.iter()).map(|(col, typ)| {
match typ {
InferredType::Integer => {
let vals: Vec<Option<i64>> = col.iter().map(|v| {
if v.is_empty() { None } else { v.replace(',', "").parse().ok() }
}).collect();
Arc::new(Int64Array::from(vals)) as ArrayRef
}
InferredType::Float => {
let vals: Vec<Option<f64>> = col.iter().map(|v| {
if v.is_empty() { None }
else { v.replace(',', "").replace('$', "").replace('%', "").parse().ok() }
}).collect();
Arc::new(Float64Array::from(vals)) as ArrayRef
}
InferredType::Boolean => {
let vals: Vec<Option<bool>> = col.iter().map(|v| {
match v.to_lowercase().as_str() {
"true" | "yes" | "1" | "y" | "t" => Some(true),
"false" | "no" | "0" | "n" | "f" => Some(false),
_ => None,
}
}).collect();
Arc::new(BooleanArray::from(vals)) as ArrayRef
}
InferredType::String => {
Arc::new(StringArray::from(col.clone())) as ArrayRef
}
}
}).collect();
let batch = RecordBatch::try_new(schema.clone(), arrays)
.map_err(|e| format!("RecordBatch error: {e}"))?;
Ok((schema, vec![batch]))
}
/// Infer column type from values. Conservative: defaults to String on ambiguity.
fn infer_column_type(values: &[String]) -> InferredType {
let non_empty: Vec<&str> = values.iter()
.map(|v| v.as_str())
.filter(|v| !v.is_empty() && *v != "NULL" && *v != "null" && *v != "N/A" && *v != "n/a")
.collect();
if non_empty.is_empty() {
return InferredType::String;
}
// Check boolean
let all_bool = non_empty.iter().all(|v| {
matches!(v.to_lowercase().as_str(), "true" | "false" | "yes" | "no" | "1" | "0" | "y" | "n" | "t" | "f")
});
if all_bool && non_empty.len() >= 2 {
// Make sure it's not just all "1" and "0" which could be integers
let has_text_bool = non_empty.iter().any(|v| {
matches!(v.to_lowercase().as_str(), "true" | "false" | "yes" | "no" | "y" | "n" | "t" | "f")
});
if has_text_bool {
return InferredType::Boolean;
}
}
// Check integer (allow commas as thousands separator)
let int_rate = non_empty.iter()
.filter(|v| v.replace(',', "").parse::<i64>().is_ok())
.count() as f64 / non_empty.len() as f64;
if int_rate > 0.95 {
return InferredType::Integer;
}
// Check float (allow $, %, commas)
let float_rate = non_empty.iter()
.filter(|v| v.replace(',', "").replace('$', "").replace('%', "").parse::<f64>().is_ok())
.count() as f64 / non_empty.len() as f64;
if float_rate > 0.95 {
return InferredType::Float;
}
InferredType::String
}
/// Sanitize column name for SQL compatibility.
fn sanitize_column_name(name: &str) -> String {
name.chars()
.map(|c| if c.is_alphanumeric() || c == '_' { c.to_ascii_lowercase() } else { '_' })
.collect::<String>()
.trim_matches('_')
.to_string()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parse_simple_csv() {
let csv = b"Name,Age,Salary\nAlice,30,50000\nBob,25,45000\n";
let (schema, batches) = parse_csv(csv).unwrap();
assert_eq!(schema.fields().len(), 3);
assert_eq!(batches[0].num_rows(), 2);
assert_eq!(schema.field(1).data_type(), &DataType::Int64);
assert_eq!(schema.field(2).data_type(), &DataType::Int64);
}
#[test]
fn parse_csv_with_mixed_types() {
let csv = b"id,value\n1,hello\n2,world\n3,N/A\n";
let (schema, _) = parse_csv(csv).unwrap();
assert_eq!(schema.field(0).data_type(), &DataType::Int64);
assert_eq!(schema.field(1).data_type(), &DataType::Utf8);
}
#[test]
fn parse_csv_with_dollar_amounts() {
let csv = b"item,price\nWidget,$29.99\nGadget,$149.50\n";
let (schema, _) = parse_csv(csv).unwrap();
assert_eq!(schema.field(1).data_type(), &DataType::Float64);
}
#[test]
fn sanitize_names() {
assert_eq!(sanitize_column_name("First Name"), "first_name");
assert_eq!(sanitize_column_name("Bill Rate ($)"), "bill_rate");
}
}
+225
View File
@@ -0,0 +1,225 @@
/// Database connector: pull tables from PostgreSQL directly into Parquet.
/// Connects, reads schema, streams rows into Arrow RecordBatches, writes Parquet.
/// No ORM, no schema definition — schema inferred from the database.
use arrow::array::{ArrayRef, BooleanArray, Float64Array, Int32Array, Int64Array, StringArray};
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use std::sync::Arc;
use tokio_postgres::{Client, NoTls, types::Type as PgType};
/// Connection config for a database.
#[derive(Debug, Clone, serde::Deserialize)]
pub struct DbConfig {
pub host: String,
pub port: u16,
pub database: String,
pub user: String,
pub password: String,
}
impl DbConfig {
pub fn connection_string(&self) -> String {
if self.password.is_empty() {
format!(
"host={} port={} dbname={} user={}",
self.host, self.port, self.database, self.user
)
} else {
format!(
"host={} port={} dbname={} user={} password={}",
self.host, self.port, self.database, self.user, self.password
)
}
}
}
/// Result of a database table import.
#[derive(Debug, Clone, serde::Serialize)]
pub struct DbIngestResult {
pub table_name: String,
pub rows: usize,
pub columns: usize,
pub schema_detected: Vec<String>, // column name: type pairs
}
/// Connect to PostgreSQL and import a table as Arrow RecordBatches.
pub async fn import_postgres_table(
config: &DbConfig,
table_name: &str,
limit: Option<usize>,
) -> Result<(Arc<Schema>, Vec<RecordBatch>, DbIngestResult), String> {
// Connect
let (client, connection) = tokio_postgres::connect(&config.connection_string(), NoTls)
.await
.map_err(|e| format!("postgres connect error: {e}"))?;
// Spawn connection handler
tokio::spawn(async move {
if let Err(e) = connection.await {
tracing::error!("postgres connection error: {e}");
}
});
tracing::info!("connected to postgres {}:{}/{}", config.host, config.port, config.database);
// Get column info
let columns = get_table_columns(&client, table_name).await?;
if columns.is_empty() {
return Err(format!("table '{}' not found or has no columns", table_name));
}
let schema_detected: Vec<String> = columns.iter()
.map(|(name, pg_type)| format!("{}: {}", name, pg_type_to_string(pg_type)))
.collect();
tracing::info!("table '{}': {} columns", table_name, columns.len());
// Build Arrow schema
let arrow_fields: Vec<Field> = columns.iter()
.map(|(name, pg_type)| Field::new(name, pg_type_to_arrow(pg_type), true))
.collect();
let schema = Arc::new(Schema::new(arrow_fields));
// Query rows
let sql = match limit {
Some(n) => format!("SELECT * FROM \"{}\" LIMIT {}", table_name, n),
None => format!("SELECT * FROM \"{}\"", table_name),
};
let rows = client.query(&sql, &[])
.await
.map_err(|e| format!("query error: {e}"))?;
let row_count = rows.len();
tracing::info!("fetched {} rows from '{}'", row_count, table_name);
if row_count == 0 {
return Ok((schema.clone(), vec![], DbIngestResult {
table_name: table_name.to_string(),
rows: 0,
columns: columns.len(),
schema_detected,
}));
}
// Convert to Arrow arrays
let mut arrays: Vec<ArrayRef> = Vec::new();
for (col_idx, (_, pg_type)) in columns.iter().enumerate() {
let array = rows_to_arrow_column(&rows, col_idx, pg_type)?;
arrays.push(array);
}
let batch = RecordBatch::try_new(schema.clone(), arrays)
.map_err(|e| format!("RecordBatch error: {e}"))?;
let result = DbIngestResult {
table_name: table_name.to_string(),
rows: row_count,
columns: columns.len(),
schema_detected,
};
Ok((schema, vec![batch], result))
}
/// List all tables in the database.
pub async fn list_postgres_tables(config: &DbConfig) -> Result<Vec<String>, String> {
let (client, connection) = tokio_postgres::connect(&config.connection_string(), NoTls)
.await
.map_err(|e| format!("connect error: {e}"))?;
tokio::spawn(async move { let _ = connection.await; });
let rows = client.query(
"SELECT table_name FROM information_schema.tables WHERE table_schema = 'public' ORDER BY table_name",
&[],
).await.map_err(|e| format!("query error: {e}"))?;
Ok(rows.iter().map(|r| r.get::<_, String>(0)).collect())
}
/// Get column names and types for a table.
async fn get_table_columns(client: &Client, table_name: &str) -> Result<Vec<(String, PgType)>, String> {
// Use a dummy query to get column types
let stmt = client.prepare(&format!("SELECT * FROM \"{}\" LIMIT 0", table_name))
.await
.map_err(|e| format!("prepare error for '{}': {e}", table_name))?;
Ok(stmt.columns().iter()
.map(|col| (col.name().to_string(), col.type_().clone()))
.collect())
}
/// Convert a Postgres type to Arrow DataType.
fn pg_type_to_arrow(pg_type: &PgType) -> DataType {
match *pg_type {
PgType::BOOL => DataType::Boolean,
PgType::INT2 | PgType::INT4 => DataType::Int32,
PgType::INT8 | PgType::OID => DataType::Int64,
PgType::FLOAT4 | PgType::FLOAT8 | PgType::NUMERIC => DataType::Float64,
_ => DataType::Utf8, // everything else → string (safe default per ADR-010)
}
}
fn pg_type_to_string(pg_type: &PgType) -> String {
format!("{}", pg_type)
}
/// Convert a column of Postgres rows to an Arrow array.
fn rows_to_arrow_column(
rows: &[tokio_postgres::Row],
col_idx: usize,
pg_type: &PgType,
) -> Result<ArrayRef, String> {
match *pg_type {
PgType::BOOL => {
let vals: Vec<Option<bool>> = rows.iter()
.map(|r| r.try_get(col_idx).ok())
.collect();
Ok(Arc::new(BooleanArray::from(vals)))
}
PgType::INT2 => {
let vals: Vec<Option<i32>> = rows.iter()
.map(|r| r.try_get::<_, i16>(col_idx).ok().map(|v| v as i32))
.collect();
Ok(Arc::new(Int32Array::from(vals)))
}
PgType::INT4 => {
let vals: Vec<Option<i32>> = rows.iter()
.map(|r| r.try_get(col_idx).ok())
.collect();
Ok(Arc::new(Int32Array::from(vals)))
}
PgType::INT8 | PgType::OID => {
let vals: Vec<Option<i64>> = rows.iter()
.map(|r| r.try_get(col_idx).ok())
.collect();
Ok(Arc::new(Int64Array::from(vals)))
}
PgType::FLOAT4 => {
let vals: Vec<Option<f64>> = rows.iter()
.map(|r| r.try_get::<_, f32>(col_idx).ok().map(|v| v as f64))
.collect();
Ok(Arc::new(Float64Array::from(vals)))
}
PgType::FLOAT8 => {
let vals: Vec<Option<f64>> = rows.iter()
.map(|r| r.try_get(col_idx).ok())
.collect();
Ok(Arc::new(Float64Array::from(vals)))
}
_ => {
// Default: try to get as string
let vals: Vec<Option<String>> = rows.iter()
.map(|r| {
r.try_get::<_, String>(col_idx).ok()
.or_else(|| r.try_get::<_, serde_json::Value>(col_idx).ok().map(|v| v.to_string()))
.or_else(|| Some("".to_string()))
})
.collect();
Ok(Arc::new(StringArray::from(vals)))
}
}
}
+103
View File
@@ -0,0 +1,103 @@
use sha2::{Digest, Sha256};
/// Detected file type from content inspection.
#[derive(Debug, Clone, PartialEq)]
pub enum FileType {
Csv,
Json,
NdJson, // newline-delimited JSON
Pdf,
Text, // plain text, SMS logs, etc.
Unknown,
}
/// Detect file type from filename extension and content sniffing.
pub fn detect_file_type(filename: &str, content: &[u8]) -> FileType {
// Extension-based first
let lower = filename.to_lowercase();
if lower.ends_with(".csv") || lower.ends_with(".tsv") {
return FileType::Csv;
}
if lower.ends_with(".json") {
// Check if it's newline-delimited JSON
if content.iter().take(4096).filter(|&&b| b == b'\n').count() > 2 {
let first_line = content.split(|&b| b == b'\n').next().unwrap_or(b"");
if first_line.starts_with(b"{") {
return FileType::NdJson;
}
}
return FileType::Json;
}
if lower.ends_with(".ndjson") || lower.ends_with(".jsonl") {
return FileType::NdJson;
}
if lower.ends_with(".pdf") {
return FileType::Pdf;
}
if lower.ends_with(".txt") || lower.ends_with(".log") || lower.ends_with(".sms") {
return FileType::Text;
}
// Content sniffing fallback
if content.starts_with(b"%PDF") {
return FileType::Pdf;
}
if content.starts_with(b"[") || content.starts_with(b"{") {
return FileType::Json;
}
// Check if it looks like CSV (has commas and newlines in first chunk)
let sample = &content[..content.len().min(4096)];
let comma_count = sample.iter().filter(|&&b| b == b',').count();
let newline_count = sample.iter().filter(|&&b| b == b'\n').count();
if comma_count > 3 && newline_count > 1 {
return FileType::Csv;
}
// If it's valid UTF-8, treat as text
if std::str::from_utf8(sample).is_ok() {
return FileType::Text;
}
FileType::Unknown
}
/// Compute SHA-256 hash of content for deduplication.
pub fn content_hash(content: &[u8]) -> String {
let mut hasher = Sha256::new();
hasher.update(content);
format!("{:x}", hasher.finalize())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn detect_csv_by_extension() {
assert_eq!(detect_file_type("data.csv", b"a,b,c\n1,2,3"), FileType::Csv);
}
#[test]
fn detect_json_by_extension() {
assert_eq!(detect_file_type("data.json", b"[{\"a\":1}]"), FileType::Json);
}
#[test]
fn detect_pdf_by_magic() {
assert_eq!(detect_file_type("unknown", b"%PDF-1.4 blah"), FileType::Pdf);
}
#[test]
fn detect_csv_by_content() {
let csv = b"name,age,city\nAlice,30,NYC\nBob,25,LA\n";
assert_eq!(detect_file_type("unknown.dat", csv), FileType::Csv);
}
#[test]
fn content_hash_deterministic() {
let h1 = content_hash(b"hello world");
let h2 = content_hash(b"hello world");
assert_eq!(h1, h2);
}
}
+167
View File
@@ -0,0 +1,167 @@
use arrow::array::{ArrayRef, Float64Array, Int64Array, StringArray, BooleanArray};
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use std::collections::BTreeMap;
use std::sync::Arc;
/// Parse JSON (array of objects or newline-delimited) into Arrow RecordBatches.
/// Nested objects are flattened: {"a": {"b": 1}} → column "a_b".
pub fn parse_json(content: &[u8]) -> Result<(Arc<Schema>, Vec<RecordBatch>), String> {
let text = std::str::from_utf8(content).map_err(|e| format!("invalid UTF-8: {e}"))?;
let text = text.trim();
// Parse into Vec<Map>
let rows: Vec<serde_json::Map<String, serde_json::Value>> = if text.starts_with('[') {
// JSON array
let arr: Vec<serde_json::Value> = serde_json::from_str(text)
.map_err(|e| format!("JSON parse error: {e}"))?;
arr.into_iter()
.filter_map(|v| v.as_object().cloned())
.collect()
} else {
// Newline-delimited JSON
text.lines()
.filter(|l| !l.trim().is_empty())
.map(|l| {
let v: serde_json::Value = serde_json::from_str(l)
.map_err(|e| format!("NDJSON line parse error: {e}"))?;
v.as_object().cloned().ok_or_else(|| "NDJSON line is not an object".into())
})
.collect::<Result<Vec<_>, String>>()?
};
if rows.is_empty() {
return Err("JSON has no records".into());
}
// Flatten and collect all columns
let mut columns: BTreeMap<String, Vec<serde_json::Value>> = BTreeMap::new();
let n_rows = rows.len();
for (row_idx, row) in rows.iter().enumerate() {
let flat = flatten_object(row, "");
// Pad existing columns that aren't in this row
for key in columns.keys().cloned().collect::<Vec<_>>() {
if !flat.contains_key(&key) {
columns.get_mut(&key).unwrap().push(serde_json::Value::Null);
}
}
// Add new columns or append values
for (key, val) in flat {
let col = columns.entry(key).or_insert_with(Vec::new);
// Pad with nulls if this column is new and we've seen prior rows
while col.len() < row_idx {
col.push(serde_json::Value::Null);
}
col.push(val);
}
}
// Pad short columns
for col in columns.values_mut() {
while col.len() < n_rows {
col.push(serde_json::Value::Null);
}
}
tracing::info!("parsed JSON: {n_rows} rows × {} columns", columns.len());
// Infer types and build schema
let mut fields = Vec::new();
let mut arrays: Vec<ArrayRef> = Vec::new();
for (name, values) in &columns {
let (dt, array) = json_column_to_arrow(values);
fields.push(Field::new(name, dt, true));
arrays.push(array);
}
let schema = Arc::new(Schema::new(fields));
let batch = RecordBatch::try_new(schema.clone(), arrays)
.map_err(|e| format!("RecordBatch error: {e}"))?;
Ok((schema, vec![batch]))
}
/// Flatten nested JSON object: {"a": {"b": 1}} → {"a_b": 1}
fn flatten_object(obj: &serde_json::Map<String, serde_json::Value>, prefix: &str) -> BTreeMap<String, serde_json::Value> {
let mut result = BTreeMap::new();
for (key, val) in obj {
let full_key = if prefix.is_empty() { key.clone() } else { format!("{prefix}_{key}") };
match val {
serde_json::Value::Object(inner) => {
result.extend(flatten_object(inner, &full_key));
}
other => {
result.insert(full_key, other.clone());
}
}
}
result
}
/// Convert a column of JSON values to an Arrow array.
fn json_column_to_arrow(values: &[serde_json::Value]) -> (DataType, ArrayRef) {
// Check if all non-null values are the same type
let non_null: Vec<&serde_json::Value> = values.iter()
.filter(|v| !v.is_null())
.collect();
if non_null.is_empty() {
return (DataType::Utf8, Arc::new(StringArray::from(vec![None::<&str>; values.len()])));
}
let all_bool = non_null.iter().all(|v| v.is_boolean());
let all_i64 = non_null.iter().all(|v| v.is_i64() || v.is_u64());
let all_f64 = non_null.iter().all(|v| v.is_number());
if all_bool {
let arr: Vec<Option<bool>> = values.iter().map(|v| v.as_bool()).collect();
(DataType::Boolean, Arc::new(BooleanArray::from(arr)))
} else if all_i64 {
let arr: Vec<Option<i64>> = values.iter().map(|v| v.as_i64()).collect();
(DataType::Int64, Arc::new(Int64Array::from(arr)))
} else if all_f64 {
let arr: Vec<Option<f64>> = values.iter().map(|v| v.as_f64()).collect();
(DataType::Float64, Arc::new(Float64Array::from(arr)))
} else {
// Default to string
let arr: Vec<Option<String>> = values.iter().map(|v| {
match v {
serde_json::Value::Null => None,
serde_json::Value::String(s) => Some(s.clone()),
other => Some(other.to_string()),
}
}).collect();
(DataType::Utf8, Arc::new(StringArray::from(arr)))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parse_json_array() {
let json = br#"[{"name":"Alice","age":30},{"name":"Bob","age":25}]"#;
let (schema, batches) = parse_json(json).unwrap();
assert_eq!(batches[0].num_rows(), 2);
assert!(schema.field_with_name("age").unwrap().data_type() == &DataType::Int64);
}
#[test]
fn parse_ndjson() {
let json = b"{\"x\":1}\n{\"x\":2}\n{\"x\":3}\n";
let (_, batches) = parse_json(json).unwrap();
assert_eq!(batches[0].num_rows(), 3);
}
#[test]
fn flatten_nested() {
let json = br#"[{"user":{"name":"Alice","address":{"city":"NYC"}},"score":9.5}]"#;
let (schema, _) = parse_json(json).unwrap();
let names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
assert!(names.contains(&"user_name"));
assert!(names.contains(&"user_address_city"));
}
}
+12
View File
@@ -0,0 +1,12 @@
pub mod db_ingest;
pub mod my_stream;
pub mod pg_stream;
pub mod schedule;
pub mod detect;
pub mod csv_ingest;
pub mod json_ingest;
pub mod pdf_ingest;
pub mod pipeline;
pub mod schema_evolution;
pub mod service;
pub mod watcher;
+410
View File
@@ -0,0 +1,410 @@
/// Streaming MySQL ingest.
///
/// Mirrors `pg_stream` for MySQL sources. Same OFFSET-paginated strategy,
/// same Parquet-streaming shape. Uses `mysql_async` (pure-rust) so we
/// don't need a C client library at build time.
///
/// Type mapping follows ADR-010 (default to string on ambiguity):
/// - Booleans (TINYINT(1)) and integer types map to Arrow Int32/Int64
/// - Floating point and decimals → Float64
/// - Everything else (text, varchar, json, date, timestamp) → Utf8
///
/// What's deliberately not supported (yet):
/// - TLS connections — `minimal` feature is plain TCP only. Upgrade
/// when a tenant actually needs it.
/// - Keyset pagination — OFFSET scans are O(N²) at multi-million-row
/// scale; same caveat as `pg_stream`.
/// - BINARY/BLOB columns — currently rendered as base64 or empty string
/// via Display-fallback.
use arrow::array::{ArrayRef, BooleanArray, Float64Array, Int32Array, Int64Array, RecordBatch, StringArray};
use arrow::datatypes::{DataType, Field, Schema};
use mysql_async::prelude::*;
use mysql_async::{Conn, Opts, Row, Value};
use parquet::arrow::ArrowWriter;
use std::sync::Arc;
/// Request shape for MySQL streaming ingest.
#[derive(Debug, Clone, serde::Deserialize)]
pub struct MyStreamRequest {
/// mysql://user:pass@host:port/db
pub dsn: String,
pub table: String,
#[serde(default)]
pub dataset_name: Option<String>,
/// Rows per fetch. Default 10_000.
#[serde(default)]
pub batch_size: Option<usize>,
/// Column to ORDER BY for stable pagination. If omitted, the first
/// column returned by the schema probe is used.
#[serde(default)]
pub order_by: Option<String>,
/// Hard cap on total rows (for sampling / previews).
#[serde(default)]
pub limit: Option<usize>,
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct MyStreamResult {
pub table: String,
pub rows: usize,
pub batches: usize,
pub columns: usize,
pub schema: Vec<String>,
pub parquet_bytes: u64,
pub duration_secs: f32,
}
/// Parsed DSN pieces. Kept local (rather than reusing pg's DbConfig) so
/// the MySQL connector doesn't depend on the PG path.
#[derive(Debug, Clone)]
pub struct MyConfig {
pub host: String,
pub port: u16,
pub user: String,
pub password: String,
pub database: String,
}
impl MyConfig {
/// Build a mysql_async `Opts` from the parsed config.
pub fn to_opts(&self) -> Opts {
let url = if self.password.is_empty() {
format!(
"mysql://{}@{}:{}/{}",
percent(&self.user), percent(&self.host), self.port, percent(&self.database),
)
} else {
format!(
"mysql://{}:{}@{}:{}/{}",
percent(&self.user), percent(&self.password),
percent(&self.host), self.port, percent(&self.database),
)
};
Opts::from_url(&url).expect("MyConfig.to_opts: rebuilt URL should parse")
}
}
/// Minimal URL-encoder for the few characters that commonly appear in
/// MySQL passwords. mysql_async's URL parser expects valid URL-encoded
/// components.
fn percent(s: &str) -> String {
let mut out = String::with_capacity(s.len());
for c in s.chars() {
match c {
'@' | ':' | '/' | '?' | '#' | '[' | ']' | ' ' | '%' => {
out.push_str(&format!("%{:02X}", c as u32));
}
_ => out.push(c),
}
}
out
}
/// Parse a mysql:// DSN.
/// Supports: mysql://[user[:password]@]host[:port][/db]
pub fn parse_dsn(dsn: &str) -> Result<MyConfig, String> {
let rest = dsn
.strip_prefix("mysql://")
.ok_or_else(|| "DSN must start with mysql://".to_string())?;
let (auth_host, database) = match rest.split_once('/') {
Some((ah, db)) => (ah, db.split('?').next().unwrap_or(db).to_string()),
None => (rest, String::new()),
};
let (userpass, hostport) = match auth_host.rsplit_once('@') {
Some((up, hp)) => (Some(up), hp),
None => (None, auth_host),
};
let (user, password) = match userpass {
Some(up) => match up.split_once(':') {
Some((u, p)) => (u.to_string(), p.to_string()),
None => (up.to_string(), String::new()),
},
None => ("root".to_string(), String::new()),
};
let (host, port) = match hostport.rsplit_once(':') {
Some((h, p)) => (
h.to_string(),
p.parse::<u16>().map_err(|_| format!("invalid port in DSN: {p}"))?,
),
None => (hostport.to_string(), 3306),
};
if host.is_empty() {
return Err("DSN has no host".into());
}
if database.is_empty() {
return Err("DSN has no database path (mysql://... /db)".into());
}
Ok(MyConfig { host, port, user, password, database })
}
/// Stream a MySQL table as Parquet bytes.
pub async fn stream_table_to_parquet(
req: &MyStreamRequest,
) -> Result<(bytes::Bytes, MyStreamResult), String> {
let t0 = std::time::Instant::now();
let config = parse_dsn(&req.dsn)?;
let batch_size = req.batch_size.unwrap_or(10_000).max(1);
let mut conn = Conn::new(config.to_opts()).await
.map_err(|e| format!("mysql connect: {e}"))?;
// Probe columns via information_schema — gives real type names that
// we can map to Arrow dtypes without needing to fetch a row first.
let columns = probe_columns(&mut conn, &config.database, &req.table).await?;
if columns.is_empty() {
return Err(format!("table '{}' not found or has no columns", req.table));
}
let arrow_fields: Vec<Field> = columns
.iter()
.map(|(name, ty)| Field::new(name, mysql_type_to_arrow(ty), true))
.collect();
let schema = Arc::new(Schema::new(arrow_fields));
let schema_report: Vec<String> = columns
.iter()
.map(|(n, t)| format!("{}:{}", n, t))
.collect();
let order_col = req.order_by.clone().unwrap_or_else(|| columns[0].0.clone());
let mut buf: Vec<u8> = Vec::with_capacity(1024 * 1024);
let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), None)
.map_err(|e| format!("ArrowWriter init: {e}"))?;
let mut total_rows: usize = 0;
let mut batch_count: usize = 0;
let row_cap = req.limit.unwrap_or(usize::MAX);
loop {
let remaining = row_cap.saturating_sub(total_rows);
if remaining == 0 { break; }
let fetch = remaining.min(batch_size);
// Backticks are MySQL's identifier quote. Forbid backticks in
// table/column names to prevent injection — neither pg nor us
// should be accepting those anyway.
if req.table.contains('`') || order_col.contains('`') {
return Err("table or order_by column contains backticks — refused".into());
}
let sql = format!(
"SELECT * FROM `{}` ORDER BY `{}` LIMIT {} OFFSET {}",
req.table, order_col, fetch, total_rows,
);
let rows: Vec<Row> = conn.query(&sql).await
.map_err(|e| format!("fetch batch at offset {total_rows}: {e}"))?;
if rows.is_empty() { break; }
let n = rows.len();
let arrays: Vec<ArrayRef> = columns
.iter()
.enumerate()
.map(|(idx, (_, ty))| rows_to_column(&rows, idx, ty))
.collect::<Result<_, _>>()?;
let batch = RecordBatch::try_new(schema.clone(), arrays)
.map_err(|e| format!("RecordBatch: {e}"))?;
writer.write(&batch).map_err(|e| format!("ArrowWriter::write: {e}"))?;
total_rows += n;
batch_count += 1;
tracing::info!(
"mysql stream '{}': fetched batch {} ({} rows, total {})",
req.table, batch_count, n, total_rows,
);
if n < fetch { break; }
}
writer.close().map_err(|e| format!("ArrowWriter::close: {e}"))?;
conn.disconnect().await.ok();
let result = MyStreamResult {
table: req.table.clone(),
rows: total_rows,
batches: batch_count,
columns: columns.len(),
schema: schema_report,
parquet_bytes: buf.len() as u64,
duration_secs: t0.elapsed().as_secs_f32(),
};
Ok((bytes::Bytes::from(buf), result))
}
async fn probe_columns(
conn: &mut Conn,
schema: &str,
table: &str,
) -> Result<Vec<(String, String)>, String> {
let sql = format!(
"SELECT COLUMN_NAME, DATA_TYPE FROM information_schema.columns \
WHERE table_schema = '{}' AND table_name = '{}' \
ORDER BY ORDINAL_POSITION",
schema.replace('\'', "''"),
table.replace('\'', "''"),
);
let rows: Vec<(String, String)> = conn.query(&sql).await
.map_err(|e| format!("probe columns: {e}"))?;
Ok(rows)
}
/// MySQL data type string → Arrow DataType. Conservative: anything we
/// don't recognize becomes Utf8 (ADR-010).
fn mysql_type_to_arrow(ty: &str) -> DataType {
match ty.to_lowercase().as_str() {
"tinyint" | "smallint" | "mediumint" | "int" | "integer" => DataType::Int32,
"bigint" => DataType::Int64,
"float" | "double" | "decimal" | "numeric" | "real" => DataType::Float64,
"bit" | "bool" | "boolean" => DataType::Boolean,
_ => DataType::Utf8,
}
}
/// Convert a single column slice of MySQL rows into an Arrow array.
fn rows_to_column(
rows: &[Row],
idx: usize,
ty: &str,
) -> Result<ArrayRef, String> {
let arrow_ty = mysql_type_to_arrow(ty);
match arrow_ty {
DataType::Boolean => {
let v: Vec<Option<bool>> = rows.iter().map(|r| cell_as_bool(r, idx)).collect();
Ok(Arc::new(BooleanArray::from(v)))
}
DataType::Int32 => {
let v: Vec<Option<i32>> = rows.iter().map(|r| cell_as_i64(r, idx).map(|n| n as i32)).collect();
Ok(Arc::new(Int32Array::from(v)))
}
DataType::Int64 => {
let v: Vec<Option<i64>> = rows.iter().map(|r| cell_as_i64(r, idx)).collect();
Ok(Arc::new(Int64Array::from(v)))
}
DataType::Float64 => {
let v: Vec<Option<f64>> = rows.iter().map(|r| cell_as_f64(r, idx)).collect();
Ok(Arc::new(Float64Array::from(v)))
}
_ => {
let v: Vec<Option<String>> = rows.iter().map(|r| cell_as_string(r, idx)).collect();
Ok(Arc::new(StringArray::from(v)))
}
}
}
fn cell(r: &Row, idx: usize) -> &Value { r.as_ref(idx).unwrap_or(&Value::NULL) }
fn cell_as_bool(r: &Row, idx: usize) -> Option<bool> {
match cell(r, idx) {
Value::NULL => None,
Value::Int(n) => Some(*n != 0),
Value::UInt(n) => Some(*n != 0),
Value::Bytes(b) => std::str::from_utf8(b).ok().and_then(|s| {
match s.to_ascii_lowercase().as_str() {
"true" | "1" | "y" | "yes" => Some(true),
"false" | "0" | "n" | "no" => Some(false),
_ => None,
}
}),
_ => None,
}
}
fn cell_as_i64(r: &Row, idx: usize) -> Option<i64> {
match cell(r, idx) {
Value::NULL => None,
Value::Int(n) => Some(*n),
Value::UInt(n) => i64::try_from(*n).ok(),
Value::Float(f) => Some(*f as i64),
Value::Double(f) => Some(*f as i64),
Value::Bytes(b) => std::str::from_utf8(b).ok().and_then(|s| s.parse().ok()),
_ => None,
}
}
fn cell_as_f64(r: &Row, idx: usize) -> Option<f64> {
match cell(r, idx) {
Value::NULL => None,
Value::Int(n) => Some(*n as f64),
Value::UInt(n) => Some(*n as f64),
Value::Float(f) => Some(*f as f64),
Value::Double(f) => Some(*f),
Value::Bytes(b) => std::str::from_utf8(b).ok().and_then(|s| s.parse().ok()),
_ => None,
}
}
fn cell_as_string(r: &Row, idx: usize) -> Option<String> {
match cell(r, idx) {
Value::NULL => None,
Value::Bytes(b) => Some(String::from_utf8_lossy(b).into_owned()),
Value::Int(n) => Some(n.to_string()),
Value::UInt(n) => Some(n.to_string()),
Value::Float(f) => Some(f.to_string()),
Value::Double(f) => Some(f.to_string()),
Value::Date(y, mo, d, h, mi, s, _us) => {
Some(format!("{:04}-{:02}-{:02} {:02}:{:02}:{:02}", y, mo, d, h, mi, s))
}
Value::Time(neg, days, h, mi, s, _us) => {
let sign = if *neg { "-" } else { "" };
Some(format!("{}{}d {:02}:{:02}:{:02}", sign, days, h, mi, s))
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parse_dsn_full() {
let c = parse_dsn("mysql://daisy:secret@db.example.com:3307/my_db").unwrap();
assert_eq!(c.host, "db.example.com");
assert_eq!(c.port, 3307);
assert_eq!(c.user, "daisy");
assert_eq!(c.password, "secret");
assert_eq!(c.database, "my_db");
}
#[test]
fn parse_dsn_default_port() {
let c = parse_dsn("mysql://root@localhost/shop").unwrap();
assert_eq!(c.port, 3306);
assert_eq!(c.user, "root");
assert_eq!(c.password, "");
}
#[test]
fn parse_dsn_no_auth() {
let c = parse_dsn("mysql://127.0.0.1:3306/analytics").unwrap();
assert_eq!(c.user, "root");
assert_eq!(c.host, "127.0.0.1");
}
#[test]
fn parse_dsn_rejects_non_mysql() {
assert!(parse_dsn("postgresql://host/db").is_err());
}
#[test]
fn parse_dsn_requires_db() {
assert!(parse_dsn("mysql://localhost:3306").is_err());
}
#[test]
fn type_map_basics() {
assert_eq!(mysql_type_to_arrow("int"), DataType::Int32);
assert_eq!(mysql_type_to_arrow("BIGINT"), DataType::Int64);
assert_eq!(mysql_type_to_arrow("decimal"), DataType::Float64);
assert_eq!(mysql_type_to_arrow("varchar"), DataType::Utf8);
assert_eq!(mysql_type_to_arrow("json"), DataType::Utf8);
assert_eq!(mysql_type_to_arrow("bool"), DataType::Boolean);
}
}
+279
View File
@@ -0,0 +1,279 @@
use arrow::array::{ArrayRef, Int32Array, StringArray};
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use std::sync::Arc;
/// Extract text from a PDF file.
///
/// Two-tier approach:
/// 1. **Text extraction (lopdf)** — fast, works on digital/native PDFs.
/// Returns one row per page: (source_file, page_number, text_content).
/// 2. **OCR fallback (Tesseract)** — if text extraction yields zero
/// pages (common with scanned documents, image-only PDFs, or PDFs
/// exported from scanners without an OCR layer), we extract embedded
/// images from the PDF and run Tesseract on each. Falls back
/// gracefully if Tesseract isn't installed (returns an error
/// directing the operator to install it).
///
/// Schema is identical for both paths so downstream consumers (chunker,
/// vectord, queryd) don't need to know how the text was produced.
pub fn parse_pdf(content: &[u8], source_filename: &str) -> Result<(Arc<Schema>, Vec<RecordBatch>), String> {
let doc = lopdf::Document::load_mem(content)
.map_err(|e| format!("PDF load error: {e}"))?;
let pages = doc.get_pages();
let mut page_numbers: Vec<i32> = Vec::new();
let mut page_texts: Vec<String> = Vec::new();
let mut sources: Vec<String> = Vec::new();
for (&page_num, _) in pages.iter() {
let text = doc.extract_text(&[page_num]).unwrap_or_default();
let text = text.trim().to_string();
if !text.is_empty() {
page_numbers.push(page_num as i32);
page_texts.push(text);
sources.push(source_filename.to_string());
}
}
// Tier 2: OCR fallback for scanned pages.
if page_numbers.is_empty() {
tracing::info!(
"no extractable text in '{}' — attempting OCR on {} pages",
source_filename, pages.len(),
);
let ocr_results = ocr_pdf_images(&doc, source_filename)?;
if ocr_results.is_empty() {
return Err(
"PDF contains no extractable text and OCR produced no output. \
Ensure tesseract is installed (`apt install tesseract-ocr tesseract-ocr-eng`)."
.into(),
);
}
for (pg, text) in ocr_results {
page_numbers.push(pg);
page_texts.push(text);
sources.push(source_filename.to_string());
}
tracing::info!(
"OCR recovered {} pages from '{}'", page_numbers.len(), source_filename,
);
} else {
tracing::info!(
"extracted {} pages with text from PDF '{}'",
page_numbers.len(), source_filename,
);
}
let schema = Arc::new(Schema::new(vec![
Field::new("source_file", DataType::Utf8, false),
Field::new("page_number", DataType::Int32, false),
Field::new("text_content", DataType::Utf8, false),
]));
let arrays: Vec<ArrayRef> = vec![
Arc::new(StringArray::from(sources)),
Arc::new(Int32Array::from(page_numbers)),
Arc::new(StringArray::from(page_texts)),
];
let batch = RecordBatch::try_new(schema.clone(), arrays)
.map_err(|e| format!("RecordBatch error: {e}"))?;
Ok((schema, vec![batch]))
}
/// Extract embedded images from the PDF, write each to a temp file,
/// run Tesseract, collect the OCR'd text. Returns (page_number, text)
/// pairs for pages that yielded non-empty output.
///
/// Image extraction strategy: walk every page's resources looking for
/// XObject streams with /Subtype /Image. Decode the raw stream bytes
/// (lopdf handles FlateDecode/DCTDecode) and write them as-is to a
/// temp file with the appropriate extension. Tesseract handles JPEG,
/// PNG, TIFF, and BMP natively.
fn ocr_pdf_images(
doc: &lopdf::Document,
source_filename: &str,
) -> Result<Vec<(i32, String)>, String> {
if !tesseract_available() {
return Err(
"tesseract not found on PATH — install with `apt install tesseract-ocr tesseract-ocr-eng`"
.into(),
);
}
let pages = doc.get_pages();
let mut results: Vec<(i32, String)> = Vec::new();
for (&page_num, &page_id) in &pages {
let images = extract_page_images(doc, page_id);
if images.is_empty() {
continue;
}
// OCR each image on this page and concatenate the text.
let mut page_text = String::new();
for (idx, img_bytes) in images.iter().enumerate() {
match ocr_image_bytes(img_bytes, page_num, idx) {
Ok(text) if !text.trim().is_empty() => {
if !page_text.is_empty() { page_text.push('\n'); }
page_text.push_str(text.trim());
}
Ok(_) => {} // empty OCR — skip
Err(e) => {
tracing::warn!(
"OCR failed on '{}' page {} image {}: {e}",
source_filename, page_num, idx,
);
}
}
}
if !page_text.is_empty() {
results.push((page_num as i32, page_text));
}
}
Ok(results)
}
/// Extract raw image bytes from a PDF page's XObject resources.
/// Handles the common case: page /Resources → /XObject dict → streams
/// with /Subtype /Image.
fn extract_page_images(
doc: &lopdf::Document,
page_id: lopdf::ObjectId,
) -> Vec<Vec<u8>> {
let mut images = Vec::new();
let page_obj = match doc.get_object(page_id) {
Ok(o) => o,
Err(_) => return images,
};
let page_dict = match page_obj.as_dict() {
Ok(d) => d,
Err(_) => return images,
};
// Resolve /Resources → /XObject
let resources = page_dict
.get(b"Resources")
.ok()
.and_then(|r| doc.dereference(r).ok())
.map(|(_, o)| o.clone());
let xobjects = resources
.as_ref()
.and_then(|r| r.as_dict().ok())
.and_then(|d| d.get(b"XObject").ok())
.and_then(|x| doc.dereference(x).ok())
.map(|(_, o)| o.clone());
let xobj_dict = match xobjects.as_ref().and_then(|x| x.as_dict().ok()) {
Some(d) => d,
None => return images,
};
for (_name, obj_ref) in xobj_dict.iter() {
let obj = match doc.dereference(obj_ref) {
Ok((_, o)) => o.clone(),
Err(_) => continue,
};
let stream = match obj.as_stream() {
Ok(s) => s,
Err(_) => continue,
};
// Check /Subtype == /Image
let is_image = stream.dict.get(b"Subtype")
.ok()
.and_then(|s| s.as_name().ok())
.map(|n| n == b"Image")
.unwrap_or(false);
if !is_image {
continue;
}
// Get decoded content — lopdf handles FlateDecode/DCTDecode.
let content = match stream.decompressed_content() {
Ok(c) => c,
Err(_) => {
// Try raw content as fallback (some DCT streams are stored raw).
stream.content.clone()
}
};
if !content.is_empty() {
images.push(content);
}
}
images
}
/// Run Tesseract on raw image bytes. Writes to a temp file, shells out,
/// reads the output. Returns the OCR'd text.
fn ocr_image_bytes(
img_bytes: &[u8],
page_num: u32,
img_idx: usize,
) -> Result<String, String> {
use std::io::Write;
use std::process::Command;
let tmp_dir = std::env::temp_dir();
// Detect image format from magic bytes to pick the right extension.
let ext = if img_bytes.starts_with(&[0xFF, 0xD8]) {
"jpg"
} else if img_bytes.starts_with(&[0x89, 0x50, 0x4E, 0x47]) {
"png"
} else if img_bytes.starts_with(b"II") || img_bytes.starts_with(b"MM") {
"tiff"
} else {
// Tesseract can often handle raw pixel data with the right options,
// but writing as PNG is safer. For truly unknown formats, try PNG
// wrapping; if Tesseract can't read it, it'll error cleanly.
"png"
};
let img_path = tmp_dir.join(format!("lh_ocr_p{page_num}_i{img_idx}.{ext}"));
let out_base = tmp_dir.join(format!("lh_ocr_p{page_num}_i{img_idx}_out"));
let out_path = tmp_dir.join(format!("lh_ocr_p{page_num}_i{img_idx}_out.txt"));
// Write image bytes to temp file.
let mut f = std::fs::File::create(&img_path)
.map_err(|e| format!("create temp image: {e}"))?;
f.write_all(img_bytes)
.map_err(|e| format!("write temp image: {e}"))?;
drop(f);
// Run Tesseract.
let status = Command::new("tesseract")
.arg(&img_path)
.arg(&out_base)
.arg("--oem").arg("3") // LSTM + legacy combined
.arg("--psm").arg("6") // assume uniform block of text
.arg("-l").arg("eng")
.stdout(std::process::Stdio::null())
.stderr(std::process::Stdio::null())
.status()
.map_err(|e| format!("tesseract exec: {e}"))?;
// Cleanup input regardless of outcome.
let _ = std::fs::remove_file(&img_path);
if !status.success() {
let _ = std::fs::remove_file(&out_path);
return Err(format!("tesseract exited with status {status}"));
}
let text = std::fs::read_to_string(&out_path).unwrap_or_default();
let _ = std::fs::remove_file(&out_path);
Ok(text)
}
/// Quick check: is `tesseract` on PATH?
fn tesseract_available() -> bool {
std::process::Command::new("tesseract")
.arg("--version")
.stdout(std::process::Stdio::null())
.stderr(std::process::Stdio::null())
.status()
.map(|s| s.success())
.unwrap_or(false)
}
+308
View File
@@ -0,0 +1,308 @@
/// Streaming PostgreSQL ingest.
///
/// The original `db_ingest::import_postgres_table` loads every row into
/// memory before emitting Parquet — fine for small tables, blows up on
/// 1M+ rows. This module paginates via `ORDER BY <pk> LIMIT N OFFSET M`
/// and streams batches into an `ArrowWriter`, closing once exhausted.
///
/// Pagination is OFFSET-based (not keyset) — simpler, works on any table
/// without needing to know the PK type. OFFSET N scans N rows per fetch,
/// which becomes quadratic for very large tables. Upgrade to keyset when
/// a user actually ingests something that hurts.
use arrow::array::{ArrayRef, BooleanArray, Float64Array, Int32Array, Int64Array, RecordBatch, StringArray};
use arrow::datatypes::{DataType, Field, Schema};
use parquet::arrow::ArrowWriter;
use std::sync::Arc;
use tokio_postgres::{Client, NoTls, types::Type as PgType};
use crate::db_ingest::DbConfig;
/// Request shape for streaming ingest — takes a DSN string and optional
/// tuning knobs.
#[derive(Debug, Clone, serde::Deserialize)]
pub struct PgStreamRequest {
/// postgresql://user:pass@host:port/db — alternative to DbConfig.
pub dsn: String,
pub table: String,
#[serde(default)]
pub dataset_name: Option<String>,
/// Rows per fetch. Default 10_000.
#[serde(default)]
pub batch_size: Option<usize>,
/// Column to ORDER BY for stable pagination. If omitted, the first
/// column returned by the schema probe is used.
#[serde(default)]
pub order_by: Option<String>,
/// Hard cap on total rows (for sampling / previews).
#[serde(default)]
pub limit: Option<usize>,
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct PgStreamResult {
pub table: String,
pub rows: usize,
pub batches: usize,
pub columns: usize,
pub schema: Vec<String>,
pub parquet_bytes: u64,
pub duration_secs: f32,
}
/// Parse a postgresql:// DSN into a DbConfig.
/// Supports: postgresql://[user[:password]@]host[:port][/db]
/// Does NOT support: query parameters (sslmode etc) — add if needed.
pub fn parse_dsn(dsn: &str) -> Result<DbConfig, String> {
let rest = dsn
.strip_prefix("postgresql://")
.or_else(|| dsn.strip_prefix("postgres://"))
.ok_or_else(|| "DSN must start with postgresql:// or postgres://".to_string())?;
// Split off path (database) first.
let (auth_host, database) = match rest.split_once('/') {
Some((ah, db)) => (ah, db.split('?').next().unwrap_or(db).to_string()),
None => (rest, "postgres".to_string()),
};
// Split user[:password] from host[:port].
let (userpass, hostport) = match auth_host.rsplit_once('@') {
Some((up, hp)) => (Some(up), hp),
None => (None, auth_host),
};
let (user, password) = match userpass {
Some(up) => match up.split_once(':') {
Some((u, p)) => (u.to_string(), p.to_string()),
None => (up.to_string(), String::new()),
},
None => ("postgres".to_string(), String::new()),
};
let (host, port) = match hostport.rsplit_once(':') {
Some((h, p)) => (
h.to_string(),
p.parse::<u16>().map_err(|_| format!("invalid port in DSN: {p}"))?,
),
None => (hostport.to_string(), 5432),
};
if host.is_empty() {
return Err("DSN has no host".into());
}
Ok(DbConfig { host, port, database, user, password })
}
/// Stream a Postgres table as Parquet bytes.
///
/// Returns the full Parquet payload + summary stats. The payload is
/// written to an in-memory buffer via `ArrowWriter` so each batch gets
/// columnar compression as it arrives — memory footprint stays roughly
/// at one batch plus Parquet footer state.
pub async fn stream_table_to_parquet(
req: &PgStreamRequest,
) -> Result<(bytes::Bytes, PgStreamResult), String> {
let t0 = std::time::Instant::now();
let config = parse_dsn(&req.dsn)?;
let batch_size = req.batch_size.unwrap_or(10_000).max(1);
let (client, connection) = tokio_postgres::connect(&config.connection_string(), NoTls)
.await
.map_err(|e| format!("postgres connect: {e}"))?;
tokio::spawn(async move {
if let Err(e) = connection.await {
tracing::error!("pg connection: {e}");
}
});
// Probe schema.
let columns = probe_columns(&client, &req.table).await?;
if columns.is_empty() {
return Err(format!("table '{}' not found or has no columns", req.table));
}
let arrow_fields: Vec<Field> = columns
.iter()
.map(|(name, pg)| Field::new(name, pg_type_to_arrow(pg), true))
.collect();
let schema = Arc::new(Schema::new(arrow_fields));
let schema_report: Vec<String> = columns
.iter()
.map(|(n, t)| format!("{}:{}", n, t))
.collect();
// Pagination key: user-specified or first column.
let order_col = req.order_by.clone().unwrap_or_else(|| columns[0].0.clone());
// Stream batches into ArrowWriter.
let mut buf: Vec<u8> = Vec::with_capacity(1024 * 1024);
let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), None)
.map_err(|e| format!("ArrowWriter init: {e}"))?;
let mut total_rows: usize = 0;
let mut batch_count: usize = 0;
let row_cap = req.limit.unwrap_or(usize::MAX);
loop {
let remaining = row_cap.saturating_sub(total_rows);
if remaining == 0 { break; }
let fetch = remaining.min(batch_size);
let sql = format!(
"SELECT * FROM \"{}\" ORDER BY \"{}\" LIMIT {} OFFSET {}",
req.table, order_col, fetch, total_rows,
);
let rows = client
.query(&sql, &[])
.await
.map_err(|e| format!("fetch batch at offset {total_rows}: {e}"))?;
if rows.is_empty() { break; }
let n = rows.len();
let arrays: Vec<ArrayRef> = columns
.iter()
.enumerate()
.map(|(idx, (_, pg))| rows_to_column(&rows, idx, pg))
.collect::<Result<_, _>>()?;
let batch = RecordBatch::try_new(schema.clone(), arrays)
.map_err(|e| format!("RecordBatch: {e}"))?;
writer.write(&batch).map_err(|e| format!("ArrowWriter::write: {e}"))?;
total_rows += n;
batch_count += 1;
tracing::info!(
"pg stream '{}': fetched batch {} ({} rows, total {})",
req.table, batch_count, n, total_rows,
);
if n < fetch { break; } // short read = end of table
}
writer.close().map_err(|e| format!("ArrowWriter::close: {e}"))?;
let parquet_bytes = buf.len() as u64;
let duration = t0.elapsed().as_secs_f32();
let result = PgStreamResult {
table: req.table.clone(),
rows: total_rows,
batches: batch_count,
columns: columns.len(),
schema: schema_report,
parquet_bytes,
duration_secs: duration,
};
Ok((bytes::Bytes::from(buf), result))
}
async fn probe_columns(client: &Client, table: &str) -> Result<Vec<(String, PgType)>, String> {
let stmt = client
.prepare(&format!("SELECT * FROM \"{}\" LIMIT 0", table))
.await
.map_err(|e| format!("prepare schema probe: {e}"))?;
Ok(stmt
.columns()
.iter()
.map(|c| (c.name().to_string(), c.type_().clone()))
.collect())
}
fn pg_type_to_arrow(pg: &PgType) -> DataType {
match *pg {
PgType::BOOL => DataType::Boolean,
PgType::INT2 | PgType::INT4 => DataType::Int32,
PgType::INT8 | PgType::OID => DataType::Int64,
PgType::FLOAT4 | PgType::FLOAT8 | PgType::NUMERIC => DataType::Float64,
_ => DataType::Utf8,
}
}
fn rows_to_column(
rows: &[tokio_postgres::Row],
idx: usize,
pg: &PgType,
) -> Result<ArrayRef, String> {
match *pg {
PgType::BOOL => {
let v: Vec<Option<bool>> = rows.iter().map(|r| r.try_get(idx).ok()).collect();
Ok(Arc::new(BooleanArray::from(v)))
}
PgType::INT2 => {
let v: Vec<Option<i32>> = rows.iter()
.map(|r| r.try_get::<_, i16>(idx).ok().map(|x| x as i32)).collect();
Ok(Arc::new(Int32Array::from(v)))
}
PgType::INT4 => {
let v: Vec<Option<i32>> = rows.iter().map(|r| r.try_get(idx).ok()).collect();
Ok(Arc::new(Int32Array::from(v)))
}
PgType::INT8 | PgType::OID => {
let v: Vec<Option<i64>> = rows.iter().map(|r| r.try_get(idx).ok()).collect();
Ok(Arc::new(Int64Array::from(v)))
}
PgType::FLOAT4 => {
let v: Vec<Option<f64>> = rows.iter()
.map(|r| r.try_get::<_, f32>(idx).ok().map(|x| x as f64)).collect();
Ok(Arc::new(Float64Array::from(v)))
}
PgType::FLOAT8 => {
let v: Vec<Option<f64>> = rows.iter().map(|r| r.try_get(idx).ok()).collect();
Ok(Arc::new(Float64Array::from(v)))
}
_ => {
// Safe-default per ADR-010: everything else becomes string.
let v: Vec<Option<String>> = rows.iter().map(|r| {
r.try_get::<_, String>(idx).ok()
.or_else(|| r.try_get::<_, serde_json::Value>(idx).ok().map(|j| j.to_string()))
.or_else(|| {
// Timestamps + uuid, rendered via Display.
r.try_get::<_, chrono::DateTime<chrono::Utc>>(idx).ok()
.map(|t| t.to_rfc3339())
})
.or_else(|| r.try_get::<_, uuid::Uuid>(idx).ok().map(|u| u.to_string()))
.or(Some(String::new()))
}).collect();
Ok(Arc::new(StringArray::from(v)))
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parse_dsn_full() {
let c = parse_dsn("postgresql://daisy:secret@db.example.com:6543/my_db").unwrap();
assert_eq!(c.host, "db.example.com");
assert_eq!(c.port, 6543);
assert_eq!(c.user, "daisy");
assert_eq!(c.password, "secret");
assert_eq!(c.database, "my_db");
}
#[test]
fn parse_dsn_minimal() {
let c = parse_dsn("postgresql://localhost/knowledge_base").unwrap();
assert_eq!(c.host, "localhost");
assert_eq!(c.port, 5432);
assert_eq!(c.user, "postgres");
assert_eq!(c.password, "");
assert_eq!(c.database, "knowledge_base");
}
#[test]
fn parse_dsn_no_password() {
let c = parse_dsn("postgres://postgres@127.0.0.1:5432/mydb").unwrap();
assert_eq!(c.user, "postgres");
assert_eq!(c.password, "");
}
#[test]
fn parse_dsn_rejects_non_pg() {
assert!(parse_dsn("http://host/db").is_err());
}
}
+219
View File
@@ -0,0 +1,219 @@
use arrow::datatypes::SchemaRef;
use arrow::record_batch::RecordBatch;
use bytes::Bytes;
use object_store::ObjectStore;
use std::sync::Arc;
use catalogd::registry::Registry;
use shared::arrow_helpers::{fingerprint_schema, record_batch_to_parquet};
use shared::types::{ObjectRef, SchemaFingerprint};
use storaged::ops;
use crate::detect::{FileType, content_hash, detect_file_type};
use crate::csv_ingest;
use crate::json_ingest;
use crate::pdf_ingest;
/// Result of an ingest operation.
#[derive(Debug, Clone, serde::Serialize)]
pub struct IngestResult {
pub dataset_name: String,
pub file_type: String,
pub rows: usize,
pub columns: usize,
pub storage_key: String,
pub content_hash: String,
pub schema_fingerprint: String,
pub deduplicated: bool,
}
/// Full ingest pipeline: detect → parse → dedup → store → register.
///
/// Bucket name (federation layer 2) is recorded on the ObjectRef so the
/// catalog knows which bucket holds the data — defaults to "primary" via
/// the ingest service's resolve_bucket helper.
pub async fn ingest_file(
filename: &str,
content: &[u8],
dataset_name: Option<&str>,
store: &Arc<dyn ObjectStore>,
registry: &Registry,
) -> Result<IngestResult, String> {
ingest_file_to_bucket(filename, content, dataset_name, "primary", store, registry).await
}
/// Same as `ingest_file` but with explicit target bucket name on the
/// ObjectRef. Header-aware ingest endpoints call this directly.
pub async fn ingest_file_to_bucket(
filename: &str,
content: &[u8],
dataset_name: Option<&str>,
bucket: &str,
store: &Arc<dyn ObjectStore>,
registry: &Registry,
) -> Result<IngestResult, String> {
// 1. Detect file type
let file_type = detect_file_type(filename, content);
tracing::info!("ingesting '{}' as {:?} ({} bytes)", filename, file_type, content.len());
// 2. Parse into Arrow
let (schema, batches) = match file_type {
FileType::Csv => csv_ingest::parse_csv(content)?,
FileType::Json | FileType::NdJson => json_ingest::parse_json(content)?,
FileType::Pdf => pdf_ingest::parse_pdf(content, filename)?,
FileType::Text => parse_text_file(content, filename)?,
FileType::Unknown => return Err(format!("unknown file type for '{filename}'")),
};
let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
let total_cols = schema.fields().len();
if total_rows == 0 {
return Err("file contains no data".into());
}
// 3. Content hash for dedup
let hash = content_hash(content);
// 4. Check if already ingested (same content hash)
let existing = registry.list().await;
if existing.iter().any(|d| d.schema_fingerprint.0 == hash) {
tracing::info!("file already ingested (hash: {}), skipping", &hash[..12]);
return Ok(IngestResult {
dataset_name: dataset_name.unwrap_or(filename).to_string(),
file_type: format!("{:?}", file_type),
rows: total_rows,
columns: total_cols,
storage_key: String::new(),
content_hash: hash,
schema_fingerprint: String::new(),
deduplicated: true,
});
}
// 5. Convert to Parquet
let mut all_parquet = Vec::new();
for batch in &batches {
let pq = record_batch_to_parquet(batch)?;
all_parquet.extend_from_slice(&pq);
}
let parquet_bytes = Bytes::from(all_parquet);
let parquet_size = parquet_bytes.len() as u64;
// 6. Store in object storage
let name = dataset_name.unwrap_or_else(|| {
filename.rsplit('/').next().unwrap_or(filename)
.rsplit('.').last().unwrap_or(filename)
});
let safe_name = sanitize_dataset_name(name);
let storage_key = format!("datasets/{}.parquet", safe_name);
ops::put(store, &storage_key, parquet_bytes).await?;
tracing::info!("stored {} as {} ({} bytes)", filename, storage_key, parquet_size);
// 7. Register in catalog with rich metadata
let schema_fp = fingerprint_schema(&schema);
let now = chrono::Utc::now();
let obj_ref = ObjectRef {
bucket: bucket.to_string(),
key: storage_key.clone(),
size_bytes: parquet_size,
created_at: now,
};
// Register base dataset
registry.register(
safe_name.clone(),
SchemaFingerprint(hash.clone()),
vec![obj_ref],
).await?;
// Auto-populate metadata: PII detection, column info, lineage, row count
let col_names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
let dataset_sensitivity = shared::pii::detect_dataset_sensitivity(&col_names);
let column_meta: Vec<shared::types::ColumnMeta> = schema.fields().iter().map(|f| {
let sens = shared::pii::detect_sensitivity(f.name());
shared::types::ColumnMeta {
name: f.name().clone(),
data_type: f.data_type().to_string(),
sensitivity: sens.clone(),
description: String::new(),
is_pii: matches!(sens, Some(shared::types::Sensitivity::Pii)),
}
}).collect();
let lineage = shared::types::Lineage {
source_system: format!("{:?}", file_type).to_lowercase(),
source_file: filename.to_string(),
ingest_job: format!("ingest-{}", now.timestamp_millis()),
ingest_timestamp: now,
parent_datasets: vec![],
};
let pii_count = column_meta.iter().filter(|c| c.is_pii).count();
if pii_count > 0 {
tracing::info!("auto-detected {} PII columns in '{}'", pii_count, safe_name);
}
let _ = registry.update_metadata(&safe_name, catalogd::registry::MetadataUpdate {
sensitivity: dataset_sensitivity,
columns: Some(column_meta),
lineage: Some(lineage),
row_count: Some(total_rows as u64),
..Default::default()
}).await;
// Phase C: if this dataset already had embeddings, they're now stale.
// mark_embeddings_stale is a no-op for never-embedded datasets.
let _ = registry.mark_embeddings_stale(&safe_name).await;
Ok(IngestResult {
dataset_name: safe_name,
file_type: format!("{:?}", file_type),
rows: total_rows,
columns: total_cols,
storage_key,
content_hash: hash,
schema_fingerprint: schema_fp.0,
deduplicated: false,
})
}
/// Parse plain text / SMS logs into a simple table.
fn parse_text_file(content: &[u8], filename: &str) -> Result<(SchemaRef, Vec<RecordBatch>), String> {
let text = String::from_utf8_lossy(content);
let lines: Vec<&str> = text.lines().filter(|l| !l.trim().is_empty()).collect();
if lines.is_empty() {
return Err("text file is empty".into());
}
let schema = Arc::new(arrow::datatypes::Schema::new(vec![
arrow::datatypes::Field::new("source_file", arrow::datatypes::DataType::Utf8, false),
arrow::datatypes::Field::new("line_number", arrow::datatypes::DataType::Int32, false),
arrow::datatypes::Field::new("text", arrow::datatypes::DataType::Utf8, false),
]));
let sources: Vec<&str> = vec![filename; lines.len()];
let line_nums: Vec<i32> = (1..=lines.len() as i32).collect();
let arrays: Vec<arrow::array::ArrayRef> = vec![
Arc::new(arrow::array::StringArray::from(sources)),
Arc::new(arrow::array::Int32Array::from(line_nums)),
Arc::new(arrow::array::StringArray::from(lines)),
];
let batch = RecordBatch::try_new(schema.clone(), arrays)
.map_err(|e| format!("RecordBatch error: {e}"))?;
Ok((schema, vec![batch]))
}
fn sanitize_dataset_name(name: &str) -> String {
let clean: String = name.chars()
.map(|c| if c.is_alphanumeric() || c == '_' { c.to_ascii_lowercase() } else { '_' })
.collect();
let trimmed = clean.trim_matches('_').to_string();
if trimmed.is_empty() { "unnamed_dataset".to_string() } else { trimmed }
}
+725
View File
@@ -0,0 +1,725 @@
//! Scheduled ingest — the production-substrate piece that turns the
//! lakehouse from "manual API toolkit" into a system that runs itself.
//!
//! A `ScheduleDef` declares "what to ingest" + "when to fire." A
//! background `Scheduler` task wakes periodically, asks each enabled
//! schedule whether it's due, and runs the matching ingest path. Outcomes
//! land back on the schedule so operators can see what happened without
//! `journalctl` archaeology.
//!
//! Storage shape: one JSON file per schedule at
//! `primary://_schedules/{id}.json`. Serializable via serde, rebuilt on
//! startup. Same write-once shape as the rest of the catalog — no DB.
//!
//! Concurrency model: single-flight per schedule_id. The scheduler holds
//! a `HashSet<String>` of currently-running ids; a tick that finds a
//! schedule already in flight skips it (no queueing — a slow ingest
//! shouldn't pile up runs). Different schedules run concurrently.
//!
//! What's deliberately not in scope here:
//! - Sub-minute cron precision. 5-field Unix cron is supported; seconds
//! are pinned to 0. Intervals cover sub-minute cases; cron is
//! easy to bolt on later.
//! - Backoff / retry policies. A failed run records the failure and
//! schedules `next_run_at` as if it succeeded — no exponential backoff.
//! - Distributed coordination. Single-machine scheduler; no leader
//! election. If you run two gateway instances they'll both fire.
use chrono::{DateTime, Duration, Utc};
use object_store::ObjectStore;
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use tokio::sync::RwLock;
use storaged::ops;
const PREFIX: &str = "_schedules";
// =================== Public types ===================
/// What kind of ingest the schedule fires.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ScheduleKind {
/// Pull a MySQL/MariaDB table via the my_stream path.
Mysql {
/// `mysql://user:pass@host:port/db` — DSN with full credentials.
/// Stored encrypted-at-rest is a future concern; for now it lives
/// in plain JSON in the catalog. Don't put production root creds
/// here without disk-level encryption.
dsn: String,
table: String,
#[serde(default)]
dataset_name: Option<String>,
#[serde(default)]
batch_size: Option<usize>,
#[serde(default)]
order_by: Option<String>,
},
/// Pull a Postgres table via the pg_stream path.
Postgres {
dsn: String,
table: String,
#[serde(default)]
dataset_name: Option<String>,
#[serde(default)]
batch_size: Option<usize>,
#[serde(default)]
order_by: Option<String>,
},
}
impl ScheduleKind {
pub fn label(&self) -> String {
match self {
ScheduleKind::Mysql { table, .. } => format!("mysql:{table}"),
ScheduleKind::Postgres { table, .. } => format!("postgres:{table}"),
}
}
}
/// When the schedule fires.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ScheduleTrigger {
/// Run every N seconds. Time anchor = compute_next_run_at decision.
Interval { secs: u64 },
/// 5-field Unix cron expression: "min hour dom month dow".
/// Seconds are implicitly 0 (minute-granularity scheduling).
/// Examples: "15 14 * * *" = 14:15 UTC daily, "0 */6 * * *" = every
/// 6 hours on the hour, "0 9 * * 1-5" = 09:00 UTC weekdays.
/// Timezone: UTC always — the scheduler's clock is UTC and we don't
/// carry a per-schedule timezone field.
Cron { expr: String },
}
/// Outcome of a single run, kept on the schedule for observability.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RunOutcome {
pub at: DateTime<Utc>,
pub success: bool,
pub message: String,
pub rows: Option<usize>,
pub duration_secs: f32,
}
/// One scheduled ingest definition.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScheduleDef {
pub id: String,
pub kind: ScheduleKind,
pub trigger: ScheduleTrigger,
pub enabled: bool,
pub created_at: DateTime<Utc>,
#[serde(default)]
pub created_by: String,
/// When the next run is due. Recomputed after each completion via
/// `compute_next_run_at`. On creation: set to "now" so the schedule
/// fires on the next scheduler tick.
pub next_run_at: DateTime<Utc>,
#[serde(default)]
pub last_run_at: Option<DateTime<Utc>>,
#[serde(default)]
pub last_outcome: Option<RunOutcome>,
/// Total successful runs and total failures since creation. Useful
/// for at-a-glance "is this schedule healthy?" checks.
#[serde(default)]
pub run_count: u64,
#[serde(default)]
pub failure_count: u64,
}
impl ScheduleDef {
/// True if this schedule should fire at `now`. Honors `enabled`.
pub fn is_due(&self, now: DateTime<Utc>) -> bool {
self.enabled && now >= self.next_run_at
}
}
// =================== Trigger semantics ===================
//
// J: THIS IS YOURS TO IMPLEMENT
// =====================================================
//
// `compute_next_run_at` decides when a schedule fires next. The current
// state is the "since last completion" semantics — simple but has real
// trade-offs vs other strategies:
//
// 1. SINCE LAST COMPLETION (current default below):
// next = completed_at + interval
// Easy to reason about. A long-running ingest delays the next fire
// by however long it took. Drift is bounded — the schedule never
// "falls behind" trying to catch up.
//
// 2. ANCHORED INTERVAL:
// next = previous_target + interval, regardless of when run finished
// Keeps wall-clock alignment ("every 30min on the :00 and :30").
// Risk: if ingest takes longer than the interval, the next run is
// "due immediately" and you can fire back-to-back forever.
//
// 3. ANCHORED + SKIP:
// Anchored, but skip a tick if the previous one is still running.
// Best of both, slightly more code.
//
// 4. CRON:
// Parse the expression, find the next match after `now`. Most
// expressive but pulls in a cron parser dep.
//
pub fn compute_next_run_at(
trigger: &ScheduleTrigger,
completed_at: DateTime<Utc>,
_previous_next: DateTime<Utc>,
) -> DateTime<Utc> {
match trigger {
ScheduleTrigger::Interval { secs } => {
completed_at + Duration::seconds(*secs as i64)
}
ScheduleTrigger::Cron { expr } => {
next_cron_fire(expr, completed_at)
// Creation-time validation (see `validate_trigger`) rejects
// unparseable expressions, so this fallback is only reached
// if the schedule JSON was hand-edited on disk. One-hour
// window keeps a bad expr from fire-looping the system.
.unwrap_or(completed_at + Duration::hours(1))
}
}
}
/// Parse a Vixie/POSIX cron expression and return the next fire time
/// strictly after `after`. Accepts 5-field (`min hour dom month dow`) and
/// 6-field (with leading `sec`) forms natively — croner handles both.
/// Day-of-week follows Unix convention: 0=Sun, 1=Mon, …, 6=Sat (7=Sun as
/// synonym). Returns None if the expression fails to parse or produces no
/// upcoming match (which can happen for year-bounded patterns in the past).
fn next_cron_fire(expr: &str, after: DateTime<Utc>) -> Option<DateTime<Utc>> {
use std::str::FromStr;
let cron = croner::Cron::from_str(expr).ok()?;
cron.find_next_occurrence(&after, false).ok()
}
/// Reject a ScheduleTrigger whose content can't be interpreted — used by
/// the HTTP handlers to fail fast at create/patch time rather than falling
/// back silently at fire time. Interval triggers get a sanity gate
/// (no zero secs); cron triggers get a full parse through croner.
pub fn validate_trigger(trigger: &ScheduleTrigger) -> Result<(), String> {
match trigger {
ScheduleTrigger::Interval { secs } => {
if *secs == 0 {
return Err("interval secs must be > 0".into());
}
Ok(())
}
ScheduleTrigger::Cron { expr } => {
use std::str::FromStr;
croner::Cron::from_str(expr)
.map_err(|e| format!("invalid cron expression '{expr}': {e}"))?;
Ok(())
}
}
}
// =================== Persistence ===================
/// Storage path for a schedule's JSON record.
fn schedule_key(id: &str) -> String {
let safe: String = id
.chars()
.map(|c| if c.is_ascii_alphanumeric() || c == '_' || c == '-' { c } else { '_' })
.collect();
format!("{PREFIX}/{safe}.json")
}
/// In-memory + on-disk schedule registry. The Scheduler reads from here;
/// HTTP CRUD writes through here.
#[derive(Clone)]
pub struct ScheduleStore {
/// Bucket where schedule JSON lives. Always primary today — a
/// future enhancement could move to per-profile schedule buckets.
store: Arc<dyn ObjectStore>,
schedules: Arc<RwLock<HashMap<String, ScheduleDef>>>,
}
impl ScheduleStore {
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
Self {
store,
schedules: Arc::new(RwLock::new(HashMap::new())),
}
}
/// Load every persisted schedule into memory. Call on startup.
pub async fn rebuild(&self) -> Result<usize, String> {
let keys = ops::list(&self.store, Some(&format!("{PREFIX}/")))
.await
.unwrap_or_default();
let mut count = 0usize;
for key in keys {
if !key.ends_with(".json") { continue; }
match ops::get(&self.store, &key).await {
Ok(bytes) => match serde_json::from_slice::<ScheduleDef>(&bytes) {
Ok(def) => {
self.schedules.write().await.insert(def.id.clone(), def);
count += 1;
}
Err(e) => tracing::warn!("schedule {key}: malformed: {e}"),
},
Err(e) => tracing::warn!("schedule {key}: read failed: {e}"),
}
}
Ok(count)
}
pub async fn list(&self) -> Vec<ScheduleDef> {
let mut out: Vec<ScheduleDef> = self.schedules.read().await.values().cloned().collect();
out.sort_by(|a, b| a.id.cmp(&b.id));
out
}
pub async fn get(&self, id: &str) -> Option<ScheduleDef> {
self.schedules.read().await.get(id).cloned()
}
pub async fn put(&self, def: ScheduleDef) -> Result<ScheduleDef, String> {
let json = serde_json::to_vec_pretty(&def).map_err(|e| e.to_string())?;
ops::put(&self.store, &schedule_key(&def.id), json.into()).await?;
self.schedules.write().await.insert(def.id.clone(), def.clone());
Ok(def)
}
pub async fn delete(&self, id: &str) -> Result<bool, String> {
// Object_store's delete is idempotent — silently OKs missing keys.
let _ = ops::delete(&self.store, &schedule_key(id)).await;
Ok(self.schedules.write().await.remove(id).is_some())
}
/// Update completion metadata + persist. Used by the scheduler
/// after a run finishes (success or failure).
pub async fn record_run(
&self,
id: &str,
outcome: RunOutcome,
) -> Result<(), String> {
let mut guard = self.schedules.write().await;
let Some(def) = guard.get_mut(id) else {
return Err(format!("schedule '{id}' not registered"));
};
def.last_run_at = Some(outcome.at);
if outcome.success {
def.run_count += 1;
} else {
def.failure_count += 1;
}
def.next_run_at = compute_next_run_at(&def.trigger, outcome.at, def.next_run_at);
def.last_outcome = Some(outcome);
let json = serde_json::to_vec_pretty(def).map_err(|e| e.to_string())?;
let key = schedule_key(id);
let store = self.store.clone();
// Drop the write lock before async I/O.
drop(guard);
ops::put(&store, &key, json.into()).await?;
Ok(())
}
}
// =================== Scheduler task ===================
/// Long-running task that fires due schedules. Spawned at gateway
/// startup. Stop semantics: drop the handle (the loop exits when its
/// store is the only reference left, or sooner if you wire a stop signal).
pub struct Scheduler {
pub store: ScheduleStore,
pub ingest: SchedulerIngestDeps,
/// Polling cadence for the loop itself. The schedules' own intervals
/// are independent — this just controls how granularly we check
/// "is anything due?" Default 10s.
pub tick_secs: u64,
}
/// What the scheduler needs to actually run an ingest. We pass in only
/// what's needed — keeps the dep set narrow.
#[derive(Clone)]
pub struct SchedulerIngestDeps {
pub store: Arc<dyn ObjectStore>,
pub registry: catalogd::registry::Registry,
pub buckets: Arc<storaged::registry::BucketRegistry>,
/// Optional Phase 16.5 hook — push a DatasetAppended event if any
/// HNSW index is bound to the dataset that just got refreshed.
pub agent_handle: Option<vectord::agent::AgentHandle>,
pub index_registry: vectord::index_registry::IndexRegistry,
}
impl Scheduler {
/// Spawn the loop in the background. Returns immediately.
pub fn spawn(self) {
tokio::spawn(async move {
self.run().await;
});
}
async fn run(self) {
tracing::info!("scheduler started — tick={}s", self.tick_secs);
let in_flight: Arc<RwLock<HashSet<String>>> = Arc::new(RwLock::new(HashSet::new()));
loop {
tokio::time::sleep(std::time::Duration::from_secs(self.tick_secs)).await;
let now = Utc::now();
let due: Vec<ScheduleDef> = self.store.list().await
.into_iter()
.filter(|s| s.is_due(now))
.collect();
if due.is_empty() { continue; }
for def in due {
{
let in_flight_guard = in_flight.read().await;
if in_flight_guard.contains(&def.id) {
// Previous run still going — skip this tick.
// We don't reschedule; next tick re-evaluates.
continue;
}
}
let store = self.store.clone();
let ingest = self.ingest.clone();
let in_flight = in_flight.clone();
let id = def.id.clone();
in_flight.write().await.insert(id.clone());
tokio::spawn(async move {
let t0 = std::time::Instant::now();
let outcome = run_schedule(&def, &ingest).await;
let elapsed = t0.elapsed().as_secs_f32();
let outcome = RunOutcome {
at: Utc::now(),
success: outcome.is_ok(),
message: match &outcome {
Ok(s) => s.clone(),
Err(e) => e.clone(),
},
rows: outcome.as_ref().ok().and_then(parse_rows_from_message),
duration_secs: elapsed,
};
if let Err(e) = store.record_run(&id, outcome).await {
tracing::warn!("scheduler: record_run({id}): {e}");
}
in_flight.write().await.remove(&id);
});
}
}
}
}
/// Pluck a "rows=N" hint out of a success message — best-effort, no
/// hard contract. Lets us surface row counts in the schedule status
/// endpoint without forcing every ingest path to return structured
/// stats here.
fn parse_rows_from_message(msg: &String) -> Option<usize> {
msg.split_whitespace()
.find_map(|tok| tok.strip_prefix("rows=").and_then(|n| n.parse().ok()))
}
/// Fire one schedule's ingest. Returns Ok(success_message) or Err(reason).
pub async fn run_schedule(
def: &ScheduleDef,
deps: &SchedulerIngestDeps,
) -> Result<String, String> {
tracing::info!("scheduler: firing '{}' ({})", def.id, def.kind.label());
match &def.kind {
ScheduleKind::Mysql { dsn, table, dataset_name, batch_size, order_by } => {
let req = crate::my_stream::MyStreamRequest {
dsn: dsn.clone(),
table: table.clone(),
dataset_name: dataset_name.clone(),
batch_size: *batch_size,
order_by: order_by.clone(),
limit: None,
};
run_mysql(req, deps).await
}
ScheduleKind::Postgres { dsn, table, dataset_name, batch_size, order_by } => {
let req = crate::pg_stream::PgStreamRequest {
dsn: dsn.clone(),
table: table.clone(),
dataset_name: dataset_name.clone(),
batch_size: *batch_size,
order_by: order_by.clone(),
limit: None,
};
run_postgres(req, deps).await
}
}
}
async fn run_mysql(
req: crate::my_stream::MyStreamRequest,
deps: &SchedulerIngestDeps,
) -> Result<String, String> {
let (parquet, result) = crate::my_stream::stream_table_to_parquet(&req).await?;
if result.rows == 0 {
return Ok(format!("rows=0 — table empty"));
}
persist_and_register(
&parquet, &result.table, req.dataset_name.as_deref(),
result.rows, "mysql", &req.dsn, deps,
).await
}
async fn run_postgres(
req: crate::pg_stream::PgStreamRequest,
deps: &SchedulerIngestDeps,
) -> Result<String, String> {
let (parquet, result) = crate::pg_stream::stream_table_to_parquet(&req).await?;
if result.rows == 0 {
return Ok(format!("rows=0 — table empty"));
}
persist_and_register(
&parquet, &result.table, req.dataset_name.as_deref(),
result.rows, "postgresql", &req.dsn, deps,
).await
}
/// Shared write+register path. Mirrors the HTTP ingest handlers in
/// service.rs so the scheduled path produces identical catalog output:
/// PII detection, lineage with redacted DSN, mark-stale + agent trigger.
async fn persist_and_register(
parquet: &bytes::Bytes,
table: &str,
dataset_name_override: Option<&str>,
rows: usize,
source_system: &str,
dsn: &str,
deps: &SchedulerIngestDeps,
) -> Result<String, String> {
use shared::types::{ColumnMeta, Lineage, ObjectRef, SchemaFingerprint, Sensitivity};
let (schema, _) = shared::arrow_helpers::parquet_to_record_batches(parquet)
.map_err(|e| format!("reparse parquet: {e}"))?;
let dataset_name = dataset_name_override.unwrap_or(table).to_string();
let storage_key = format!("datasets/{}.parquet", dataset_name);
let parquet_size = parquet.len() as u64;
let bucket = "primary".to_string();
let target_store = deps.buckets.get(&bucket)?;
ops::put(&target_store, &storage_key, parquet.clone()).await?;
let schema_fp = shared::arrow_helpers::fingerprint_schema(&schema);
let now = Utc::now();
deps.registry.register(
dataset_name.clone(),
SchemaFingerprint(schema_fp.0),
vec![ObjectRef {
bucket: bucket.clone(),
key: storage_key.clone(),
size_bytes: parquet_size,
created_at: now,
}],
).await?;
let col_names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
let sensitivity = shared::pii::detect_dataset_sensitivity(&col_names);
let columns: Vec<ColumnMeta> = schema.fields().iter().map(|f| {
let sens = shared::pii::detect_sensitivity(f.name());
ColumnMeta {
name: f.name().clone(),
data_type: f.data_type().to_string(),
sensitivity: sens.clone(),
description: String::new(),
is_pii: matches!(sens, Some(Sensitivity::Pii)),
}
}).collect();
let lineage = Lineage {
source_system: source_system.to_string(),
source_file: format!("dsn: {}", redact_dsn(dsn)),
ingest_job: format!("scheduled-{}-{}", source_system, now.timestamp_millis()),
ingest_timestamp: now,
parent_datasets: vec![],
};
let _ = deps.registry.update_metadata(&dataset_name, catalogd::registry::MetadataUpdate {
sensitivity,
columns: Some(columns),
lineage: Some(lineage),
row_count: Some(rows as u64),
..Default::default()
}).await;
let _ = deps.registry.mark_embeddings_stale(&dataset_name).await;
if let Some(agent) = &deps.agent_handle {
let bound = deps.index_registry.list(Some(&dataset_name), None).await;
for meta in bound {
let event = vectord::agent::TriggerEvent::dataset_appended(
meta.index_name.clone(), &dataset_name,
);
let _ = agent.enqueue(event).await;
}
}
Ok(format!("rows={rows} dataset={dataset_name} bytes={parquet_size}"))
}
/// Same generic redactor used in service.rs — duplicated here to avoid
/// pulling service::redact_dsn behind a `pub` visibility just for this.
fn redact_dsn(dsn: &str) -> String {
let scheme_end = match dsn.find("://") {
Some(i) => i + 3,
None => return dsn.to_string(),
};
let at_idx = match dsn.rfind('@') {
Some(i) if i > scheme_end => i,
_ => return dsn.to_string(),
};
let userpass = &dsn[scheme_end..at_idx];
let colon_offset = match userpass.find(':') {
Some(i) => i,
None => return dsn.to_string(),
};
let colon_idx = scheme_end + colon_offset;
format!("{}:***{}", &dsn[..colon_idx], &dsn[at_idx..])
}
// =================== Tests ===================
#[cfg(test)]
mod tests {
use super::*;
use chrono::Timelike;
fn mk(secs: u64) -> ScheduleDef {
ScheduleDef {
id: "test".into(),
kind: ScheduleKind::Mysql {
dsn: "mysql://x@host/db".into(),
table: "t".into(),
dataset_name: None,
batch_size: None,
order_by: None,
},
trigger: ScheduleTrigger::Interval { secs },
enabled: true,
created_at: Utc::now(),
created_by: "tester".into(),
next_run_at: Utc::now(),
last_run_at: None,
last_outcome: None,
run_count: 0,
failure_count: 0,
}
}
#[test]
fn is_due_respects_enabled() {
let mut s = mk(60);
s.next_run_at = Utc::now() - Duration::seconds(5);
assert!(s.is_due(Utc::now()));
s.enabled = false;
assert!(!s.is_due(Utc::now()));
}
#[test]
fn is_due_respects_future_time() {
let mut s = mk(60);
s.next_run_at = Utc::now() + Duration::seconds(60);
assert!(!s.is_due(Utc::now()));
}
#[test]
fn next_run_since_last_completion_advances() {
let s = mk(120);
let completed = Utc::now();
let next = compute_next_run_at(&s.trigger, completed, completed);
let delta = next - completed;
assert_eq!(delta.num_seconds(), 120);
}
#[test]
fn cron_every_minute_fires_within_sixty_seconds() {
let trig = ScheduleTrigger::Cron { expr: "* * * * *".into() };
let now = Utc::now();
let next = compute_next_run_at(&trig, now, now);
let delta = next - now;
assert!(delta.num_seconds() > 0 && delta.num_seconds() <= 60,
"expected next fire within 60s, got {}s", delta.num_seconds());
}
#[test]
fn cron_daily_at_1415_utc_is_within_24h() {
// 14:15 UTC daily — whether we're before or after 14:15 today, the
// next fire is at most 24h out.
let trig = ScheduleTrigger::Cron { expr: "15 14 * * *".into() };
let now: DateTime<Utc> = "2026-04-20T10:00:00Z".parse().unwrap();
let next = compute_next_run_at(&trig, now, now);
assert!(next > now);
assert!((next - now).num_hours() <= 24);
// And the minute-of-hour is 15.
assert_eq!(next.minute(), 15);
assert_eq!(next.hour(), 14);
}
#[test]
fn cron_weekday_skips_weekend() {
// 09:00 UTC weekdays only. On a Saturday, next fire is Monday.
let trig = ScheduleTrigger::Cron { expr: "0 9 * * 1-5".into() };
// 2026-04-18 is a Saturday.
let sat: DateTime<Utc> = "2026-04-18T10:00:00Z".parse().unwrap();
let next = compute_next_run_at(&trig, sat, sat);
// Monday 2026-04-20 at 09:00 UTC.
assert_eq!(next, "2026-04-20T09:00:00Z".parse::<DateTime<Utc>>().unwrap());
}
#[test]
fn cron_six_field_seconds_granularity() {
// 6-field (seconds included): fire every minute at 30s past the minute.
let trig = ScheduleTrigger::Cron { expr: "30 * * * * *".into() };
let now = Utc::now();
let next = compute_next_run_at(&trig, now, now);
assert_eq!(next.second(), 30);
}
#[test]
fn validate_rejects_bad_cron() {
let bad = ScheduleTrigger::Cron { expr: "not a cron".into() };
let err = validate_trigger(&bad).unwrap_err();
assert!(err.contains("cron"), "error should mention cron, got: {err}");
}
#[test]
fn validate_rejects_wrong_field_count() {
// 4 fields — neither 5 nor 6.
let bad = ScheduleTrigger::Cron { expr: "* * * *".into() };
assert!(validate_trigger(&bad).is_err());
}
#[test]
fn validate_rejects_zero_interval() {
let bad = ScheduleTrigger::Interval { secs: 0 };
assert!(validate_trigger(&bad).is_err());
}
#[test]
fn validate_accepts_good_cron() {
let good = ScheduleTrigger::Cron { expr: "0 */6 * * *".into() };
assert!(validate_trigger(&good).is_ok());
}
#[test]
fn schedule_kind_label() {
let k = ScheduleKind::Mysql {
dsn: "x".into(), table: "customers".into(),
dataset_name: None, batch_size: None, order_by: None,
};
assert_eq!(k.label(), "mysql:customers");
}
#[test]
fn redact_dsn_handles_mysql_and_postgres() {
assert_eq!(redact_dsn("mysql://u:secret@h/db"), "mysql://u:***@h/db");
assert_eq!(redact_dsn("postgresql://u:secret@h:5432/db"), "postgresql://u:***@h:5432/db");
assert_eq!(redact_dsn("mysql://u@h/db"), "mysql://u@h/db"); // no password
}
}
+279
View File
@@ -0,0 +1,279 @@
/// Schema evolution: detect changes between ingests and suggest migration rules.
/// When a source changes format (columns renamed, added, removed, type changed),
/// the system detects the diff and can auto-map using AI or heuristic matching.
use arrow::datatypes::{DataType, SchemaRef};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
/// A detected change between two schema versions.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SchemaChange {
pub change_type: ChangeType,
pub column: String,
pub detail: String,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum ChangeType {
Added, // new column in incoming schema
Removed, // column in old schema not in new
TypeChanged, // same name, different type
Renamed, // likely renamed (fuzzy match)
}
/// A migration rule — how to map old schema to new.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MigrationRule {
pub source_column: String,
pub target_column: String,
pub transform: String, // "direct", "cast", "concat", "split", "drop"
pub confidence: f32, // 0.0 - 1.0
pub auto_generated: bool,
}
/// Compare two schemas and detect changes.
pub fn diff_schemas(old: &SchemaRef, new: &SchemaRef) -> Vec<SchemaChange> {
let mut changes = Vec::new();
let old_fields: HashMap<&str, &DataType> = old.fields().iter()
.map(|f| (f.name().as_str(), f.data_type()))
.collect();
let new_fields: HashMap<&str, &DataType> = new.fields().iter()
.map(|f| (f.name().as_str(), f.data_type()))
.collect();
// Check for removed columns
for (name, _dt) in &old_fields {
if !new_fields.contains_key(name) {
// Check if it was renamed (fuzzy match)
let possible_rename = find_similar_column(name, &new_fields, &old_fields);
if let Some((new_name, confidence)) = possible_rename {
changes.push(SchemaChange {
change_type: ChangeType::Renamed,
column: name.to_string(),
detail: format!("likely renamed to '{}' (confidence: {:.0}%)", new_name, confidence * 100.0),
});
} else {
changes.push(SchemaChange {
change_type: ChangeType::Removed,
column: name.to_string(),
detail: "column no longer present in incoming data".into(),
});
}
}
}
// Check for added or type-changed columns
for (name, new_dt) in &new_fields {
match old_fields.get(name) {
None => {
// Skip if we already flagged this as a rename target
let is_rename_target = changes.iter().any(|c| {
c.change_type == ChangeType::Renamed && c.detail.contains(name)
});
if !is_rename_target {
changes.push(SchemaChange {
change_type: ChangeType::Added,
column: name.to_string(),
detail: format!("new column (type: {})", new_dt),
});
}
}
Some(old_dt) => {
if old_dt != new_dt {
changes.push(SchemaChange {
change_type: ChangeType::TypeChanged,
column: name.to_string(),
detail: format!("{} → {}", old_dt, new_dt),
});
}
}
}
}
changes
}
/// Generate migration rules from detected changes.
pub fn generate_rules(changes: &[SchemaChange]) -> Vec<MigrationRule> {
changes.iter().filter_map(|change| {
match change.change_type {
ChangeType::Renamed => {
// Extract the new name from the detail string
let new_name = change.detail
.split('\'')
.nth(1)
.unwrap_or(&change.column);
Some(MigrationRule {
source_column: change.column.clone(),
target_column: new_name.to_string(),
transform: "direct".into(),
confidence: 0.8,
auto_generated: true,
})
}
ChangeType::TypeChanged => {
Some(MigrationRule {
source_column: change.column.clone(),
target_column: change.column.clone(),
transform: "cast".into(),
confidence: 0.9,
auto_generated: true,
})
}
ChangeType::Added => {
Some(MigrationRule {
source_column: String::new(),
target_column: change.column.clone(),
transform: "new_column".into(),
confidence: 1.0,
auto_generated: true,
})
}
ChangeType::Removed => {
Some(MigrationRule {
source_column: change.column.clone(),
target_column: String::new(),
transform: "drop".into(),
confidence: 0.7,
auto_generated: true,
})
}
}
}).collect()
}
/// Build an AI prompt to suggest migration rules for complex cases.
pub fn build_migration_prompt(
old_schema: &SchemaRef,
new_schema: &SchemaRef,
changes: &[SchemaChange],
) -> String {
let old_cols: Vec<String> = old_schema.fields().iter()
.map(|f| format!(" {} ({})", f.name(), f.data_type()))
.collect();
let new_cols: Vec<String> = new_schema.fields().iter()
.map(|f| format!(" {} ({})", f.name(), f.data_type()))
.collect();
let change_list: Vec<String> = changes.iter()
.map(|c| format!(" [{:?}] {} — {}", c.change_type, c.column, c.detail))
.collect();
format!(
"You are a data migration expert. A data source has changed its schema.\n\n\
OLD SCHEMA:\n{}\n\n\
NEW SCHEMA:\n{}\n\n\
DETECTED CHANGES:\n{}\n\n\
For each change, suggest a migration rule as JSON:\n\
{{\"source_column\": \"old_name\", \"target_column\": \"new_name\", \"transform\": \"direct|cast|concat|split|drop\", \"confidence\": 0.0-1.0}}\n\n\
Output ONLY a JSON array of rules, nothing else.",
old_cols.join("\n"),
new_cols.join("\n"),
change_list.join("\n"),
)
}
/// Fuzzy match: find a column in new schema that's similar to a removed old column.
fn find_similar_column<'a>(
old_name: &str,
new_fields: &HashMap<&'a str, &DataType>,
old_fields: &HashMap<&str, &DataType>,
) -> Option<(String, f32)> {
let old_lower = old_name.to_lowercase();
let old_parts: Vec<&str> = old_lower.split('_').collect();
let mut best_match: Option<(String, f32)> = None;
for (new_name, _) in new_fields {
// Skip columns that already exist in old schema (not renamed)
if old_fields.contains_key(new_name) { continue; }
let new_lower = new_name.to_lowercase();
let new_parts: Vec<&str> = new_lower.split('_').collect();
// Exact substring match
if old_lower.contains(&new_lower) || new_lower.contains(&old_lower) {
let confidence = 0.85;
if best_match.as_ref().map_or(true, |(_, c)| confidence > *c) {
best_match = Some((new_name.to_string(), confidence));
}
continue;
}
// Shared word parts (e.g., "first_name" → "fname", "full_name")
let shared: usize = old_parts.iter()
.filter(|p| new_parts.iter().any(|np| np.contains(*p) || p.contains(np)))
.count();
if shared > 0 {
let confidence = shared as f32 / old_parts.len().max(new_parts.len()) as f32 * 0.75;
if confidence >= 0.3 && best_match.as_ref().map_or(true, |(_, c)| confidence > *c) {
best_match = Some((new_name.to_string(), confidence));
}
}
}
best_match
}
#[cfg(test)]
mod tests {
use super::*;
use arrow::datatypes::{Field, Schema};
use std::sync::Arc;
fn schema(fields: Vec<(&str, DataType)>) -> SchemaRef {
Arc::new(Schema::new(
fields.into_iter().map(|(n, t)| Field::new(n, t, true)).collect::<Vec<_>>()
))
}
#[test]
fn detect_added_column() {
let old = schema(vec![("id", DataType::Int64), ("name", DataType::Utf8)]);
let new = schema(vec![("id", DataType::Int64), ("name", DataType::Utf8), ("email", DataType::Utf8)]);
let changes = diff_schemas(&old, &new);
assert_eq!(changes.len(), 1);
assert_eq!(changes[0].change_type, ChangeType::Added);
assert_eq!(changes[0].column, "email");
}
#[test]
fn detect_removed_column() {
let old = schema(vec![("id", DataType::Int64), ("name", DataType::Utf8), ("age", DataType::Int64)]);
let new = schema(vec![("id", DataType::Int64), ("name", DataType::Utf8)]);
let changes = diff_schemas(&old, &new);
assert!(changes.iter().any(|c| c.column == "age" && c.change_type == ChangeType::Removed));
}
#[test]
fn detect_type_change() {
let old = schema(vec![("id", DataType::Int64), ("score", DataType::Int64)]);
let new = schema(vec![("id", DataType::Int64), ("score", DataType::Float64)]);
let changes = diff_schemas(&old, &new);
assert_eq!(changes.len(), 1);
assert_eq!(changes[0].change_type, ChangeType::TypeChanged);
}
#[test]
fn detect_rename() {
let old = schema(vec![("first_name", DataType::Utf8), ("last_name", DataType::Utf8)]);
let new = schema(vec![("full_name", DataType::Utf8)]);
let changes = diff_schemas(&old, &new);
// Should detect first_name → full_name as a potential rename
assert!(changes.iter().any(|c| c.change_type == ChangeType::Renamed));
}
#[test]
fn generate_rules_from_changes() {
let old = schema(vec![("id", DataType::Int64), ("name", DataType::Utf8)]);
let new = schema(vec![("id", DataType::Int64), ("name", DataType::Utf8), ("email", DataType::Utf8)]);
let changes = diff_schemas(&old, &new);
let rules = generate_rules(&changes);
assert_eq!(rules.len(), 1);
assert_eq!(rules[0].transform, "new_column");
}
}
+758
View File
@@ -0,0 +1,758 @@
use axum::{
Json, Router,
extract::{Multipart, Path, Query, State},
http::{HeaderMap, StatusCode},
response::IntoResponse,
routing::{get, post},
};
use bytes::Bytes;
use object_store::ObjectStore;
use serde::Deserialize;
use std::sync::Arc;
use catalogd::registry::Registry;
use crate::{db_ingest, my_stream, pg_stream, pipeline, schedule};
use shared::arrow_helpers::record_batch_to_parquet;
use shared::types::{ObjectRef, SchemaFingerprint};
use storaged::ops;
use storaged::registry::BucketRegistry;
#[derive(Clone)]
pub struct IngestState {
pub store: Arc<dyn ObjectStore>,
pub registry: Registry,
/// Federation layer 2: lookup target bucket from request headers.
pub buckets: Arc<BucketRegistry>,
/// Phase 16.5: when ingest marks a dataset's embeddings stale, push
/// a `DatasetAppended` trigger so the autotune agent can schedule
/// re-trials. Holds the agent handle (no-op when agent disabled).
pub agent_handle: vectord::agent::AgentHandle,
/// Used to look up which HNSW indexes are attached to the
/// just-ingested dataset. Each matching index gets one trigger.
pub index_registry: vectord::index_registry::IndexRegistry,
/// Scheduled-ingest registry. The scheduler task runs against this
/// store; HTTP CRUD endpoints write through it.
pub schedules: schedule::ScheduleStore,
/// Event journal for ADR-012 mutation history. Optional for back-compat
/// with callers (like scheduled ingest tests) that don't wire it yet.
/// When present, successful ingests emit a record_ingest event — closes
/// P9-001 on the file-upload path. (2026-04-23)
pub journal: Option<journald::journal::Journal>,
}
/// Push `DatasetAppended` triggers for every HNSW index bound to this
/// dataset. Called after a successful ingest. Logs failures but never
/// fails the ingest — the agent is best-effort, not load-bearing.
async fn notify_agent_on_append(state: &IngestState, dataset_name: &str) {
let indexes = state.index_registry.list(Some(dataset_name), None).await;
if indexes.is_empty() {
return;
}
for meta in indexes {
let event = vectord::agent::TriggerEvent::dataset_appended(
meta.index_name.clone(), dataset_name,
);
if let Err(e) = state.agent_handle.enqueue(event).await {
tracing::warn!("agent enqueue failed for '{}': {}", meta.index_name, e);
} else {
tracing::info!(
"agent: enqueued DatasetAppended({}) for index '{}'",
dataset_name, meta.index_name,
);
}
}
}
/// Resolve the target bucket from `X-Lakehouse-Bucket` header.
/// Returns `(bucket_name, store_for_writes)`. Falls back to "primary"
/// when the header is absent. Returns Err with the canonical bucket
/// list when the header names an unknown bucket.
fn resolve_bucket(
headers: &HeaderMap,
buckets: &BucketRegistry,
) -> Result<(String, Arc<dyn ObjectStore>), (StatusCode, String)> {
let target = headers
.get("x-lakehouse-bucket")
.and_then(|h| h.to_str().ok())
.map(|s| s.to_string())
.unwrap_or_else(|| buckets.default_name().to_string());
match buckets.get(&target) {
Ok(store) => Ok((target, store)),
Err(_) => Err((
StatusCode::NOT_FOUND,
format!("unknown bucket '{}' — use GET /storage/buckets to list", target),
)),
}
}
pub fn router(state: IngestState) -> Router {
Router::new()
.route("/health", get(health))
.route("/file", post(ingest_file))
.route("/postgres/tables", post(list_pg_tables))
.route("/postgres/import", post(import_pg_table))
.route("/db", post(ingest_db_stream))
.route("/mysql", post(ingest_mysql_stream))
// Phase E: scheduled ingest
.route("/schedules", get(list_schedules).post(create_schedule))
.route("/schedules/{id}", get(get_schedule)
.patch(patch_schedule)
.delete(delete_schedule))
.route("/schedules/{id}/run-now", post(run_schedule_now))
.with_state(state)
}
async fn health() -> &'static str {
"ingestd ok"
}
#[derive(Deserialize)]
struct IngestQuery {
name: Option<String>,
}
async fn ingest_file(
State(state): State<IngestState>,
Query(query): Query<IngestQuery>,
headers: HeaderMap,
mut multipart: Multipart,
) -> impl IntoResponse {
let (bucket, store) = resolve_bucket(&headers, &state.buckets)?;
let field = match multipart.next_field().await {
Ok(Some(f)) => f,
Ok(None) => return Err((StatusCode::BAD_REQUEST, "no file uploaded".to_string())),
Err(e) => return Err((StatusCode::BAD_REQUEST, format!("multipart error: {e}"))),
};
let filename = field.file_name().unwrap_or("unknown").to_string();
let content = field.bytes().await
.map_err(|e| (StatusCode::BAD_REQUEST, format!("read error: {e}")))?;
tracing::info!(
"ingest '{}' ({} bytes) -> bucket={}",
filename, content.len(), bucket,
);
let dataset_name = query.name.as_deref();
match pipeline::ingest_file_to_bucket(&filename, &content, dataset_name, &bucket, &store, &state.registry).await {
Ok(result) => {
if !result.deduplicated {
notify_agent_on_append(&state, &result.dataset_name).await;
// P9-001 fix (2026-04-23): emit a mutation event on every
// non-deduplicated ingest. Dedup no-ops don't need events
// (ADR-020 register() is already idempotent on same fingerprint).
if let Some(ref journal) = state.journal {
if let Err(e) = journal.record_ingest(
&result.dataset_name,
result.rows as usize,
"ingest_api",
&filename,
).await {
tracing::warn!(
"journal record_ingest failed for '{}': {}",
result.dataset_name, e,
);
}
}
}
if result.deduplicated {
Ok((StatusCode::OK, Json(result)))
} else {
Ok((StatusCode::CREATED, Json(result)))
}
}
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
// --- PostgreSQL Import ---
/// List tables in a PostgreSQL database.
async fn list_pg_tables(
Json(config): Json<db_ingest::DbConfig>,
) -> impl IntoResponse {
match db_ingest::list_postgres_tables(&config).await {
Ok(tables) => Ok(Json(tables)),
Err(e) => Err((StatusCode::BAD_GATEWAY, e)),
}
}
#[derive(Deserialize)]
struct PgImportRequest {
#[serde(flatten)]
config: db_ingest::DbConfig,
table: String,
/// Override dataset name (defaults to table name)
dataset_name: Option<String>,
/// Max rows to import (None = all)
limit: Option<usize>,
}
/// Import a PostgreSQL table into the lakehouse.
async fn import_pg_table(
State(state): State<IngestState>,
Json(req): Json<PgImportRequest>,
) -> Result<(StatusCode, Json<serde_json::Value>), (StatusCode, String)> {
tracing::info!("importing postgres table '{}' from {}:{}/{}",
req.table, req.config.host, req.config.port, req.config.database);
// Import from Postgres
let (schema, batches, db_result) = db_ingest::import_postgres_table(
&req.config, &req.table, req.limit,
).await.map_err(|e| (StatusCode::BAD_GATEWAY, e))?;
if batches.is_empty() || db_result.rows == 0 {
return Ok((StatusCode::OK, Json(serde_json::json!({
"table": req.table,
"rows": 0,
"message": "table is empty",
}))));
}
// Convert to Parquet
let mut all_parquet = Vec::new();
for batch in &batches {
let pq = record_batch_to_parquet(batch)
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
all_parquet.extend_from_slice(&pq);
}
let dataset_name = req.dataset_name.unwrap_or_else(|| req.table.clone());
let storage_key = format!("datasets/{}.parquet", dataset_name);
let parquet_size = all_parquet.len() as u64;
// Store
ops::put(&state.store, &storage_key, Bytes::from(all_parquet))
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
// Register
let schema_fp = shared::arrow_helpers::fingerprint_schema(&schema);
let now = chrono::Utc::now();
state.registry.register(
dataset_name.clone(),
SchemaFingerprint(schema_fp.0),
vec![ObjectRef {
bucket: "data".to_string(),
key: storage_key.clone(),
size_bytes: parquet_size,
created_at: now,
}],
).await.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
// Auto-populate metadata (PII detection, lineage)
let col_names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
let sensitivity = shared::pii::detect_dataset_sensitivity(&col_names);
let columns: Vec<shared::types::ColumnMeta> = schema.fields().iter().map(|f| {
let sens = shared::pii::detect_sensitivity(f.name());
shared::types::ColumnMeta {
name: f.name().clone(),
data_type: f.data_type().to_string(),
sensitivity: sens.clone(),
description: String::new(),
is_pii: matches!(sens, Some(shared::types::Sensitivity::Pii)),
}
}).collect();
let lineage = shared::types::Lineage {
source_system: "postgresql".to_string(),
source_file: format!("{}:{}/{}.{}", req.config.host, req.config.port, req.config.database, req.table),
ingest_job: format!("pg-import-{}", now.timestamp_millis()),
ingest_timestamp: now,
parent_datasets: vec![],
};
let _ = state.registry.update_metadata(&dataset_name, catalogd::registry::MetadataUpdate {
sensitivity,
columns: Some(columns),
lineage: Some(lineage),
row_count: Some(db_result.rows as u64),
..Default::default()
}).await;
tracing::info!("imported '{}' from postgres: {} rows → {}", dataset_name, db_result.rows, storage_key);
Ok((StatusCode::CREATED, Json(serde_json::json!({
"dataset_name": dataset_name,
"table": req.table,
"rows": db_result.rows,
"columns": db_result.columns,
"schema": db_result.schema_detected,
"storage_key": storage_key,
"size_bytes": parquet_size,
}))))
}
/// Streaming DSN-based ingest. Paginates via ORDER BY + LIMIT/OFFSET so
/// large tables don't blow up memory. This is the path the task spec calls
/// `POST /ingest/db` with `{dsn, table, dataset_name}`.
async fn ingest_db_stream(
State(state): State<IngestState>,
headers: HeaderMap,
Json(req): Json<pg_stream::PgStreamRequest>,
) -> Result<(StatusCode, Json<serde_json::Value>), (StatusCode, String)> {
let (bucket, store) = resolve_bucket(&headers, &state.buckets)?;
tracing::info!(
"pg stream ingest: table='{}' dataset='{:?}' bucket='{}' batch_size={:?}",
req.table, req.dataset_name, bucket, req.batch_size,
);
// Stream from Postgres into Parquet bytes.
let (parquet, stream_result) = pg_stream::stream_table_to_parquet(&req)
.await
.map_err(|e| (StatusCode::BAD_GATEWAY, e))?;
if stream_result.rows == 0 {
return Ok((StatusCode::OK, Json(serde_json::json!({
"table": req.table,
"rows": 0,
"message": "table is empty",
}))));
}
// Recover schema from the parquet footer — keeps a single source of truth
// for the schema fingerprint + column metadata.
let (schema, _) = shared::arrow_helpers::parquet_to_record_batches(&parquet)
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("reparse parquet: {e}")))?;
let dataset_name = req.dataset_name.clone().unwrap_or_else(|| req.table.clone());
let storage_key = format!("datasets/{}.parquet", dataset_name);
let size_bytes = parquet.len() as u64;
ops::put(&store, &storage_key, parquet)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
let schema_fp = shared::arrow_helpers::fingerprint_schema(&schema);
let now = chrono::Utc::now();
state.registry.register(
dataset_name.clone(),
SchemaFingerprint(schema_fp.0),
vec![ObjectRef {
bucket: bucket.clone(),
key: storage_key.clone(),
size_bytes,
created_at: now,
}],
).await.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
// Rich metadata: PII auto-detect, lineage, row_count, columns.
let col_names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
let sensitivity = shared::pii::detect_dataset_sensitivity(&col_names);
let columns: Vec<shared::types::ColumnMeta> = schema.fields().iter().map(|f| {
let sens = shared::pii::detect_sensitivity(f.name());
shared::types::ColumnMeta {
name: f.name().clone(),
data_type: f.data_type().to_string(),
sensitivity: sens.clone(),
description: String::new(),
is_pii: matches!(sens, Some(shared::types::Sensitivity::Pii)),
}
}).collect();
let lineage = shared::types::Lineage {
source_system: "postgresql".to_string(),
source_file: format!("dsn: {}", redact_dsn(&req.dsn)),
ingest_job: format!("pg-stream-{}", now.timestamp_millis()),
ingest_timestamp: now,
parent_datasets: vec![],
};
let _ = state.registry.update_metadata(&dataset_name, catalogd::registry::MetadataUpdate {
sensitivity,
columns: Some(columns),
lineage: Some(lineage),
row_count: Some(stream_result.rows as u64),
..Default::default()
}).await;
// Phase C: mark embeddings stale if the dataset already had a vector
// index. No-op for newly-created datasets.
let _ = state.registry.mark_embeddings_stale(&dataset_name).await;
// Phase 16.5: notify autotune agent. No-op if no indexes attached.
notify_agent_on_append(&state, &dataset_name).await;
Ok((StatusCode::CREATED, Json(serde_json::json!({
"dataset_name": dataset_name,
"table": stream_result.table,
"rows": stream_result.rows,
"batches": stream_result.batches,
"columns": stream_result.columns,
"schema": stream_result.schema,
"storage_key": storage_key,
"size_bytes": size_bytes,
"duration_secs": stream_result.duration_secs,
}))))
}
/// Streaming MySQL ingest. Mirrors `ingest_db_stream` — same pagination +
/// Parquet-footer schema recovery + PII/lineage metadata. Different driver,
/// otherwise identical.
async fn ingest_mysql_stream(
State(state): State<IngestState>,
headers: HeaderMap,
Json(req): Json<my_stream::MyStreamRequest>,
) -> Result<(StatusCode, Json<serde_json::Value>), (StatusCode, String)> {
let (bucket, store) = resolve_bucket(&headers, &state.buckets)?;
tracing::info!(
"mysql stream ingest: table='{}' dataset='{:?}' bucket='{}' batch_size={:?}",
req.table, req.dataset_name, bucket, req.batch_size,
);
let (parquet, stream_result) = my_stream::stream_table_to_parquet(&req)
.await
.map_err(|e| (StatusCode::BAD_GATEWAY, e))?;
if stream_result.rows == 0 {
return Ok((StatusCode::OK, Json(serde_json::json!({
"table": req.table,
"rows": 0,
"message": "table is empty",
}))));
}
let (schema, _) = shared::arrow_helpers::parquet_to_record_batches(&parquet)
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("reparse parquet: {e}")))?;
let dataset_name = req.dataset_name.clone().unwrap_or_else(|| req.table.clone());
let storage_key = format!("datasets/{}.parquet", dataset_name);
let size_bytes = parquet.len() as u64;
ops::put(&store, &storage_key, parquet)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
let schema_fp = shared::arrow_helpers::fingerprint_schema(&schema);
let now = chrono::Utc::now();
state.registry.register(
dataset_name.clone(),
SchemaFingerprint(schema_fp.0),
vec![ObjectRef {
bucket: bucket.clone(),
key: storage_key.clone(),
size_bytes,
created_at: now,
}],
).await.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
let col_names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
let sensitivity = shared::pii::detect_dataset_sensitivity(&col_names);
let columns: Vec<shared::types::ColumnMeta> = schema.fields().iter().map(|f| {
let sens = shared::pii::detect_sensitivity(f.name());
shared::types::ColumnMeta {
name: f.name().clone(),
data_type: f.data_type().to_string(),
sensitivity: sens.clone(),
description: String::new(),
is_pii: matches!(sens, Some(shared::types::Sensitivity::Pii)),
}
}).collect();
let lineage = shared::types::Lineage {
source_system: "mysql".to_string(),
source_file: format!("dsn: {}", redact_dsn(&req.dsn)),
ingest_job: format!("mysql-stream-{}", now.timestamp_millis()),
ingest_timestamp: now,
parent_datasets: vec![],
};
let _ = state.registry.update_metadata(&dataset_name, catalogd::registry::MetadataUpdate {
sensitivity,
columns: Some(columns),
lineage: Some(lineage),
row_count: Some(stream_result.rows as u64),
..Default::default()
}).await;
// Phase C: mark stale + Phase 16.5: notify agent.
let _ = state.registry.mark_embeddings_stale(&dataset_name).await;
notify_agent_on_append(&state, &dataset_name).await;
Ok((StatusCode::CREATED, Json(serde_json::json!({
"dataset_name": dataset_name,
"table": stream_result.table,
"rows": stream_result.rows,
"batches": stream_result.batches,
"columns": stream_result.columns,
"schema": stream_result.schema,
"storage_key": storage_key,
"size_bytes": size_bytes,
"duration_secs": stream_result.duration_secs,
}))))
}
/// Redact the password in any `scheme://user:pass@host/...` DSN. Works
/// for postgres, mysql, and any other scheme using the same shape.
fn redact_dsn(dsn: &str) -> String {
let scheme_end = match dsn.find("://") {
Some(i) => i + 3,
None => return dsn.to_string(),
};
let at_idx = match dsn.rfind('@') {
Some(i) if i > scheme_end => i,
_ => return dsn.to_string(),
};
// Find the password separator (first `:` after scheme, before `@`).
let userpass = &dsn[scheme_end..at_idx];
let colon_offset = match userpass.find(':') {
Some(i) => i,
None => return dsn.to_string(), // no password in DSN
};
let colon_idx = scheme_end + colon_offset;
format!("{}:***{}", &dsn[..colon_idx], &dsn[at_idx..])
}
// --- Phase E: Scheduled ingest CRUD endpoints ---
#[derive(Deserialize)]
struct CreateScheduleRequest {
/// Stable id used in URLs and on disk. Required so callers can
/// PATCH/DELETE deterministically. Sanitized on storage but the
/// caller's id is what gets persisted as the canonical key.
id: String,
kind: schedule::ScheduleKind,
trigger: schedule::ScheduleTrigger,
#[serde(default = "default_enabled")]
enabled: bool,
#[serde(default)]
created_by: String,
}
fn default_enabled() -> bool { true }
async fn create_schedule(
State(state): State<IngestState>,
Json(req): Json<CreateScheduleRequest>,
) -> impl IntoResponse {
if state.schedules.get(&req.id).await.is_some() {
return Err((StatusCode::CONFLICT, format!("schedule '{}' already exists", req.id)));
}
if let Err(e) = schedule::validate_trigger(&req.trigger) {
return Err((StatusCode::BAD_REQUEST, e));
}
let now = chrono::Utc::now();
let def = schedule::ScheduleDef {
id: req.id,
kind: req.kind,
trigger: req.trigger,
enabled: req.enabled,
created_at: now,
created_by: req.created_by,
// Fire on the next tick — operators usually want to verify
// their schedule works without waiting for the full interval.
next_run_at: now,
last_run_at: None,
last_outcome: None,
run_count: 0,
failure_count: 0,
};
match state.schedules.put(def).await {
Ok(d) => Ok((StatusCode::CREATED, Json(d))),
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
async fn list_schedules(State(state): State<IngestState>) -> impl IntoResponse {
Json(state.schedules.list().await)
}
async fn get_schedule(
State(state): State<IngestState>,
Path(id): Path<String>,
) -> impl IntoResponse {
match state.schedules.get(&id).await {
Some(d) => Ok(Json(d)),
None => Err((StatusCode::NOT_FOUND, format!("schedule '{id}' not found"))),
}
}
#[derive(Deserialize)]
struct PatchScheduleRequest {
/// Toggle on/off without changing anything else.
#[serde(default)]
enabled: Option<bool>,
/// Replace the trigger spec — useful for bumping intervals.
#[serde(default)]
trigger: Option<schedule::ScheduleTrigger>,
}
async fn patch_schedule(
State(state): State<IngestState>,
Path(id): Path<String>,
Json(req): Json<PatchScheduleRequest>,
) -> impl IntoResponse {
let Some(mut def) = state.schedules.get(&id).await else {
return Err((StatusCode::NOT_FOUND, format!("schedule '{id}' not found")));
};
if let Some(t) = &req.trigger {
if let Err(e) = schedule::validate_trigger(t) {
return Err((StatusCode::BAD_REQUEST, e));
}
}
if let Some(e) = req.enabled { def.enabled = e; }
if let Some(t) = req.trigger { def.trigger = t; }
match state.schedules.put(def).await {
Ok(d) => Ok(Json(d)),
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
async fn delete_schedule(
State(state): State<IngestState>,
Path(id): Path<String>,
) -> impl IntoResponse {
match state.schedules.delete(&id).await {
Ok(true) => Ok(Json(serde_json::json!({ "deleted": id }))),
Ok(false) => Err((StatusCode::NOT_FOUND, format!("schedule '{id}' not found"))),
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
/// Force a schedule to fire right now, regardless of its next_run_at.
/// Records the outcome and reschedules. Useful for "test my new schedule
/// works without waiting an hour."
async fn run_schedule_now(
State(state): State<IngestState>,
Path(id): Path<String>,
) -> impl IntoResponse {
let Some(def) = state.schedules.get(&id).await else {
return Err((StatusCode::NOT_FOUND, format!("schedule '{id}' not found")));
};
let deps = schedule::SchedulerIngestDeps {
store: state.store.clone(),
registry: state.registry.clone(),
buckets: state.buckets.clone(),
agent_handle: Some(state.agent_handle.clone()),
index_registry: state.index_registry.clone(),
};
let t0 = std::time::Instant::now();
let result = schedule::run_schedule(&def, &deps).await;
let elapsed = t0.elapsed().as_secs_f32();
let outcome = schedule::RunOutcome {
at: chrono::Utc::now(),
success: result.is_ok(),
message: match &result {
Ok(s) => s.clone(),
Err(e) => e.clone(),
},
rows: result.as_ref().ok().and_then(|m: &String| {
m.split_whitespace()
.find_map(|tok| tok.strip_prefix("rows=").and_then(|n| n.parse().ok()))
}),
duration_secs: elapsed,
};
if let Err(e) = state.schedules.record_run(&id, outcome.clone()).await {
tracing::warn!("schedule {id}: record_run failed: {e}");
}
Ok(Json(outcome))
}
// ─── Tests ───
#[cfg(test)]
mod journal_integration_tests {
//! P9-001 integration test: prove that a successful ingest produces a
//! journal.record_ingest event. Block 2 on PR #10 was "journal event
//! verified live" being unbacked by the diff. This test makes the
//! verification committed and reproducible.
use journald::journal::{Event, Journal};
use object_store::memory::InMemory;
use std::sync::Arc;
// Helper: build a bare Journal against an in-memory object store.
// Flush threshold 1 so every recorded event is persisted immediately.
fn test_journal() -> Journal {
let store: Arc<dyn object_store::ObjectStore> = Arc::new(InMemory::new());
Journal::new(store, 1)
}
#[tokio::test]
async fn journal_record_ingest_increments_counter() {
// Arrange — fresh journal, counter starts at zero.
let journal = test_journal();
let stats0 = journal.stats().await;
assert_eq!(stats0.total_events_created, 0);
assert_eq!(stats0.buffer_events, 0);
// Act — simulate what the /ingest/file success path does.
journal
.record_ingest("test_dataset", 42, "ingest_api", "probe.csv")
.await
.expect("record_ingest should succeed");
// Assert — counter advanced, event exists. With threshold=1 the
// event flushed to store; with threshold>N it would be in-buffer.
let stats1 = journal.stats().await;
assert_eq!(stats1.total_events_created, 1, "counter should reflect one recorded event");
// Assert — the event is retrievable by entity.
let history = journal
.get_entity_history("batch:42")
.await
.expect("history lookup");
assert_eq!(history.len(), 1, "one event should be visible in history");
let ev = &history[0];
assert_eq!(ev.action, "ingest");
assert_eq!(ev.entity_type, "test_dataset");
assert_eq!(ev.actor, "ingest_api");
assert!(
ev.new_value.contains("probe.csv"),
"new_value should carry source filename, got: {}",
ev.new_value
);
}
#[tokio::test]
async fn optional_journal_field_none_is_valid_back_compat() {
// IngestState.journal is Option<Journal>. Back-compat path: when
// the field is None, the ingest handler MUST still succeed — the
// journal call is fire-and-forget, never load-bearing.
//
// This test asserts the type shape: Option<Journal> is what we
// expect. If a refactor makes it mandatory, this test forces an
// explicit re-consideration.
let none_journal: Option<Journal> = None;
assert!(none_journal.is_none());
let some_journal: Option<Journal> = Some(test_journal());
assert!(some_journal.is_some());
}
#[tokio::test]
async fn journal_record_event_fields_match_adr_012_schema() {
// ADR-012 locks the event schema: entity_type, entity_id, field,
// action, old_value, new_value, actor, source, workspace_id plus
// the auto-assigned event_id + timestamp. This test pins the
// field names so a future refactor can't silently drop one.
let journal = test_journal();
let base = Event {
event_id: String::new(),
timestamp: chrono::Utc::now(),
entity_type: "candidate".into(),
entity_id: "CAND-0001".into(),
field: "phone".into(),
action: "update".into(),
old_value: "555-0000".into(),
new_value: "555-9999".into(),
actor: "recruiter".into(),
source: "api".into(),
workspace_id: "ws-x".into(),
};
journal.record(base).await.expect("record should accept full-schema event");
let h = journal
.get_entity_history("CAND-0001")
.await
.expect("lookup");
assert_eq!(h.len(), 1);
assert_eq!(h[0].field, "phone");
assert_eq!(h[0].old_value, "555-0000");
assert_eq!(h[0].new_value, "555-9999");
assert_eq!(h[0].workspace_id, "ws-x");
}
}
+124
View File
@@ -0,0 +1,124 @@
/// File watcher: monitors a directory for new files and auto-ingests them.
/// Drop a CSV, JSON, PDF, or text file → auto-detected → Parquet → queryable.
/// Processed files are moved to a "processed/" subdirectory.
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;
use tokio::time;
use catalogd::registry::Registry;
use object_store::ObjectStore;
use crate::pipeline;
/// Watch config.
pub struct WatchConfig {
pub watch_dir: PathBuf,
pub processed_dir: PathBuf,
pub failed_dir: PathBuf,
pub poll_interval: Duration,
}
impl Default for WatchConfig {
fn default() -> Self {
Self {
watch_dir: PathBuf::from("./inbox"),
processed_dir: PathBuf::from("./inbox/processed"),
failed_dir: PathBuf::from("./inbox/failed"),
poll_interval: Duration::from_secs(10),
}
}
}
/// Start the file watcher as a background task.
pub fn spawn_watcher(
config: WatchConfig,
store: Arc<dyn ObjectStore>,
registry: Registry,
) -> tokio::task::JoinHandle<()> {
// Create directories
std::fs::create_dir_all(&config.watch_dir).ok();
std::fs::create_dir_all(&config.processed_dir).ok();
std::fs::create_dir_all(&config.failed_dir).ok();
tracing::info!("file watcher started: watching {}", config.watch_dir.display());
tokio::spawn(async move {
let mut interval = time::interval(config.poll_interval);
loop {
interval.tick().await;
if let Err(e) = process_inbox(&config, &store, &registry).await {
tracing::error!("watcher error: {e}");
}
}
})
}
/// Scan inbox directory and process any new files.
async fn process_inbox(
config: &WatchConfig,
store: &Arc<dyn ObjectStore>,
registry: &Registry,
) -> Result<(), String> {
let entries = std::fs::read_dir(&config.watch_dir)
.map_err(|e| format!("can't read inbox: {e}"))?;
for entry in entries {
let entry = match entry {
Ok(e) => e,
Err(_) => continue,
};
let path = entry.path();
// Skip directories and hidden files. Bind filename once via
// let-else so the subsequent use is unwrap-free — previous
// version relied on a map_or guard above + an .unwrap() here
// being consistent, which is a fragile invariant.
if path.is_dir() { continue; }
let Some(fn_os) = path.file_name() else { continue; };
let filename = fn_os.to_string_lossy().to_string();
if filename.starts_with('.') { continue; }
tracing::info!("watcher: found new file '{}'", filename);
// Read file
let content = match std::fs::read(&path) {
Ok(c) => c,
Err(e) => {
tracing::error!("watcher: can't read '{}': {e}", filename);
move_file(&path, &config.failed_dir.join(&filename));
continue;
}
};
// Ingest
match pipeline::ingest_file(&filename, &content, None, store, registry).await {
Ok(result) => {
if result.deduplicated {
tracing::info!("watcher: '{}' already ingested, skipping", filename);
} else {
tracing::info!("watcher: ingested '{}' → {} ({} rows, {} cols)",
filename, result.dataset_name, result.rows, result.columns);
}
// Move to processed
move_file(&path, &config.processed_dir.join(&filename));
}
Err(e) => {
tracing::error!("watcher: failed to ingest '{}': {e}", filename);
move_file(&path, &config.failed_dir.join(&filename));
}
}
}
Ok(())
}
fn move_file(from: &Path, to: &Path) {
if let Err(e) = std::fs::rename(from, to) {
// rename fails across filesystems, try copy+delete
if let Err(e2) = std::fs::copy(from, to).and_then(|_| std::fs::remove_file(from)) {
tracing::warn!("can't move {} → {}: {e}, copy failed: {e2}", from.display(), to.display());
}
}
}
+18
View File
@@ -0,0 +1,18 @@
[package]
name = "journald"
version = "0.1.0"
edition = "2024"
[dependencies]
shared = { path = "../shared" }
storaged = { path = "../storaged" }
tokio = { workspace = true }
axum = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }
tracing = { workspace = true }
arrow = { workspace = true }
parquet = { workspace = true }
bytes = { workspace = true }
chrono = { workspace = true }
object_store = { workspace = true }
+322
View File
@@ -0,0 +1,322 @@
/// Append-only event journal.
/// Every mutation to any dataset is recorded here as an immutable event.
/// Events are never modified or deleted — this IS the audit trail.
///
/// Storage: events buffer in memory, flush to Parquet periodically.
/// Query: load Parquet files, filter by entity/field/actor/time.
use arrow::array::{ArrayRef, RecordBatch, StringArray};
use arrow::datatypes::{DataType, Field, Schema};
use chrono::{DateTime, Utc};
use object_store::ObjectStore;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tokio::sync::RwLock;
use shared::arrow_helpers::record_batch_to_parquet;
use storaged::ops;
/// A single mutation event. Immutable once created.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Event {
pub event_id: String,
pub timestamp: DateTime<Utc>,
pub entity_type: String, // "candidate", "placement", "client", etc.
pub entity_id: String, // "CAND-00342"
pub field: String, // "phone", "*" for insert/delete
pub action: String, // "insert", "update", "delete"
pub old_value: String, // "" for inserts
pub new_value: String, // "" for deletes
pub actor: String, // "Sarah", "ingest_pipeline", "auto_matcher"
pub source: String, // "api", "ingest", "workspace", "agent"
pub workspace_id: String, // "" if not workspace-scoped
}
/// Arrow schema for journal Parquet files.
fn journal_schema() -> Arc<Schema> {
Arc::new(Schema::new(vec![
Field::new("event_id", DataType::Utf8, false),
Field::new("timestamp", DataType::Utf8, false),
Field::new("entity_type", DataType::Utf8, false),
Field::new("entity_id", DataType::Utf8, false),
Field::new("field", DataType::Utf8, false),
Field::new("action", DataType::Utf8, false),
Field::new("old_value", DataType::Utf8, true),
Field::new("new_value", DataType::Utf8, true),
Field::new("actor", DataType::Utf8, false),
Field::new("source", DataType::Utf8, false),
Field::new("workspace_id", DataType::Utf8, true),
]))
}
/// The event journal — buffers events in memory, flushes to Parquet.
#[derive(Clone)]
pub struct Journal {
buffer: Arc<RwLock<Vec<Event>>>,
store: Arc<dyn ObjectStore>,
flush_threshold: usize,
event_counter: Arc<std::sync::atomic::AtomicU64>,
}
impl Journal {
pub fn new(store: Arc<dyn ObjectStore>, flush_threshold: usize) -> Self {
Self {
buffer: Arc::new(RwLock::new(Vec::new())),
store,
flush_threshold,
event_counter: Arc::new(std::sync::atomic::AtomicU64::new(0)),
}
}
/// Record a single event. Triggers flush if buffer exceeds threshold.
pub async fn record(&self, mut event: Event) -> Result<(), String> {
let seq = self.event_counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
event.event_id = format!("evt-{}-{seq}", event.timestamp.timestamp_millis());
let should_flush = {
let mut buf = self.buffer.write().await;
buf.push(event);
buf.len() >= self.flush_threshold
};
if should_flush {
self.flush().await?;
}
Ok(())
}
/// Record a batch insert (multiple rows added to a dataset).
pub async fn record_insert(
&self,
entity_type: &str,
entity_id: &str,
actor: &str,
source: &str,
workspace_id: &str,
) -> Result<(), String> {
self.record(Event {
event_id: String::new(),
timestamp: Utc::now(),
entity_type: entity_type.to_string(),
entity_id: entity_id.to_string(),
field: "*".to_string(),
action: "insert".to_string(),
old_value: String::new(),
new_value: String::new(),
actor: actor.to_string(),
source: source.to_string(),
workspace_id: workspace_id.to_string(),
}).await
}
/// Record a field update.
pub async fn record_update(
&self,
entity_type: &str,
entity_id: &str,
field: &str,
old_value: &str,
new_value: &str,
actor: &str,
source: &str,
workspace_id: &str,
) -> Result<(), String> {
self.record(Event {
event_id: String::new(),
timestamp: Utc::now(),
entity_type: entity_type.to_string(),
entity_id: entity_id.to_string(),
field: field.to_string(),
action: "update".to_string(),
old_value: old_value.to_string(),
new_value: new_value.to_string(),
actor: actor.to_string(),
source: source.to_string(),
workspace_id: workspace_id.to_string(),
}).await
}
/// Record a batch of events from a dataset ingest.
pub async fn record_ingest(
&self,
entity_type: &str,
count: usize,
actor: &str,
source_file: &str,
) -> Result<(), String> {
self.record(Event {
event_id: String::new(),
timestamp: Utc::now(),
entity_type: entity_type.to_string(),
entity_id: format!("batch:{count}"),
field: "*".to_string(),
action: "ingest".to_string(),
old_value: String::new(),
new_value: format!("{count} rows from {source_file}"),
actor: actor.to_string(),
source: "ingest".to_string(),
workspace_id: String::new(),
}).await
}
/// Flush buffer to Parquet. Called automatically at threshold or manually.
pub async fn flush(&self) -> Result<usize, String> {
let events = {
let mut buf = self.buffer.write().await;
if buf.is_empty() { return Ok(0); }
std::mem::take(&mut *buf)
};
let n = events.len();
let batch = events_to_batch(&events)?;
let parquet = record_batch_to_parquet(&batch)?;
let ts = Utc::now().format("%Y%m%d_%H%M%S_%3f");
let key = format!("journal/{ts}.parquet");
ops::put(&self.store, &key, parquet).await?;
tracing::info!("journal: flushed {n} events to {key}");
Ok(n)
}
/// Query journal: get all events for a specific entity.
pub async fn get_entity_history(&self, entity_id: &str) -> Result<Vec<Event>, String> {
let mut all_events = Vec::new();
// Check in-memory buffer first
{
let buf = self.buffer.read().await;
for e in buf.iter() {
if e.entity_id == entity_id {
all_events.push(e.clone());
}
}
}
// Load from persisted journal files
let keys = ops::list(&self.store, Some("journal/")).await?;
for key in &keys {
if !key.ends_with(".parquet") { continue; }
let data = ops::get(&self.store, key).await?;
let (_, batches) = shared::arrow_helpers::parquet_to_record_batches(&data)?;
for batch in &batches {
let events = batch_to_events(batch)?;
for e in events {
if e.entity_id == entity_id {
all_events.push(e);
}
}
}
}
all_events.sort_by_key(|e| e.timestamp);
Ok(all_events)
}
/// Query journal: get recent events (last N).
pub async fn get_recent(&self, limit: usize) -> Result<Vec<Event>, String> {
let mut all_events = Vec::new();
// Buffer
{
let buf = self.buffer.read().await;
all_events.extend(buf.iter().cloned());
}
// Persisted (load most recent files first)
let mut keys = ops::list(&self.store, Some("journal/")).await?;
keys.sort();
keys.reverse();
for key in &keys {
if !key.ends_with(".parquet") { continue; }
if all_events.len() >= limit { break; }
let data = ops::get(&self.store, key).await?;
let (_, batches) = shared::arrow_helpers::parquet_to_record_batches(&data)?;
for batch in &batches {
all_events.extend(batch_to_events(batch)?);
}
}
all_events.sort_by_key(|e| e.timestamp);
all_events.reverse();
all_events.truncate(limit);
Ok(all_events)
}
/// Get buffer stats.
pub async fn stats(&self) -> JournalStats {
let buf = self.buffer.read().await;
let keys = ops::list(&self.store, Some("journal/")).await.unwrap_or_default();
JournalStats {
buffer_events: buf.len(),
flush_threshold: self.flush_threshold,
persisted_files: keys.len(),
total_events_created: self.event_counter.load(std::sync::atomic::Ordering::Relaxed),
}
}
}
#[derive(Debug, Clone, Serialize)]
pub struct JournalStats {
pub buffer_events: usize,
pub flush_threshold: usize,
pub persisted_files: usize,
pub total_events_created: u64,
}
fn events_to_batch(events: &[Event]) -> Result<RecordBatch, String> {
let schema = journal_schema();
let arrays: Vec<ArrayRef> = vec![
Arc::new(StringArray::from(events.iter().map(|e| e.event_id.as_str()).collect::<Vec<_>>())),
Arc::new(StringArray::from(events.iter().map(|e| e.timestamp.to_rfc3339()).collect::<Vec<_>>())),
Arc::new(StringArray::from(events.iter().map(|e| e.entity_type.as_str()).collect::<Vec<_>>())),
Arc::new(StringArray::from(events.iter().map(|e| e.entity_id.as_str()).collect::<Vec<_>>())),
Arc::new(StringArray::from(events.iter().map(|e| e.field.as_str()).collect::<Vec<_>>())),
Arc::new(StringArray::from(events.iter().map(|e| e.action.as_str()).collect::<Vec<_>>())),
Arc::new(StringArray::from(events.iter().map(|e| e.old_value.as_str()).collect::<Vec<_>>())),
Arc::new(StringArray::from(events.iter().map(|e| e.new_value.as_str()).collect::<Vec<_>>())),
Arc::new(StringArray::from(events.iter().map(|e| e.actor.as_str()).collect::<Vec<_>>())),
Arc::new(StringArray::from(events.iter().map(|e| e.source.as_str()).collect::<Vec<_>>())),
Arc::new(StringArray::from(events.iter().map(|e| e.workspace_id.as_str()).collect::<Vec<_>>())),
];
RecordBatch::try_new(schema, arrays).map_err(|e| format!("batch error: {e}"))
}
fn batch_to_events(batch: &RecordBatch) -> Result<Vec<Event>, String> {
let get_col = |i: usize| -> Result<&StringArray, String> {
batch.column(i).as_any().downcast_ref::<StringArray>()
.ok_or_else(|| format!("column {i} not string"))
};
let event_ids = get_col(0)?;
let timestamps = get_col(1)?;
let entity_types = get_col(2)?;
let entity_ids = get_col(3)?;
let fields = get_col(4)?;
let actions = get_col(5)?;
let old_values = get_col(6)?;
let new_values = get_col(7)?;
let actors = get_col(8)?;
let sources = get_col(9)?;
let workspace_ids = get_col(10)?;
let mut events = Vec::with_capacity(batch.num_rows());
for i in 0..batch.num_rows() {
events.push(Event {
event_id: event_ids.value(i).to_string(),
timestamp: timestamps.value(i).parse().unwrap_or_else(|_| Utc::now()),
entity_type: entity_types.value(i).to_string(),
entity_id: entity_ids.value(i).to_string(),
field: fields.value(i).to_string(),
action: actions.value(i).to_string(),
old_value: old_values.value(i).to_string(),
new_value: new_values.value(i).to_string(),
actor: actors.value(i).to_string(),
source: sources.value(i).to_string(),
workspace_id: workspace_ids.value(i).to_string(),
});
}
Ok(events)
}
+2
View File
@@ -0,0 +1,2 @@
pub mod journal;
pub mod service;
+131
View File
@@ -0,0 +1,131 @@
use axum::{
Json, Router,
extract::{Path, Query, State},
http::StatusCode,
response::IntoResponse,
routing::{get, post},
};
use serde::Deserialize;
use crate::journal::{Event, Journal};
pub fn router(journal: Journal) -> Router {
Router::new()
.route("/health", get(health))
.route("/event", post(record_event))
.route("/update", post(record_update))
.route("/flush", post(flush))
.route("/history/{entity_id}", get(entity_history))
.route("/recent", get(recent_events))
.route("/stats", get(stats))
.with_state(journal)
}
async fn health() -> &'static str {
"journald ok"
}
// --- Record events ---
#[derive(Deserialize)]
struct RecordEventRequest {
entity_type: String,
entity_id: String,
action: String,
field: Option<String>,
old_value: Option<String>,
new_value: Option<String>,
actor: String,
source: String,
workspace_id: Option<String>,
}
async fn record_event(
State(journal): State<Journal>,
Json(req): Json<RecordEventRequest>,
) -> impl IntoResponse {
let event = Event {
event_id: String::new(), // assigned by journal
timestamp: chrono::Utc::now(),
entity_type: req.entity_type,
entity_id: req.entity_id,
field: req.field.unwrap_or_else(|| "*".to_string()),
action: req.action,
old_value: req.old_value.unwrap_or_default(),
new_value: req.new_value.unwrap_or_default(),
actor: req.actor,
source: req.source,
workspace_id: req.workspace_id.unwrap_or_default(),
};
match journal.record(event).await {
Ok(()) => Ok((StatusCode::CREATED, "event recorded")),
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
#[derive(Deserialize)]
struct UpdateRequest {
entity_type: String,
entity_id: String,
field: String,
old_value: String,
new_value: String,
actor: String,
workspace_id: Option<String>,
}
async fn record_update(
State(journal): State<Journal>,
Json(req): Json<UpdateRequest>,
) -> impl IntoResponse {
match journal.record_update(
&req.entity_type, &req.entity_id, &req.field,
&req.old_value, &req.new_value, &req.actor, "api",
&req.workspace_id.unwrap_or_default(),
).await {
Ok(()) => Ok((StatusCode::CREATED, "update recorded")),
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
// --- Flush ---
async fn flush(State(journal): State<Journal>) -> impl IntoResponse {
match journal.flush().await {
Ok(n) => Ok(format!("flushed {n} events")),
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
// --- Query ---
async fn entity_history(
State(journal): State<Journal>,
Path(entity_id): Path<String>,
) -> impl IntoResponse {
match journal.get_entity_history(&entity_id).await {
Ok(events) => Ok(Json(events)),
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
#[derive(Deserialize)]
struct RecentQuery {
limit: Option<usize>,
}
async fn recent_events(
State(journal): State<Journal>,
Query(q): Query<RecentQuery>,
) -> impl IntoResponse {
let limit = q.limit.unwrap_or(50);
match journal.get_recent(limit).await {
Ok(events) => Ok(Json(events)),
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
async fn stats(State(journal): State<Journal>) -> impl IntoResponse {
Json(journal.stats().await)
}
+42
View File
@@ -0,0 +1,42 @@
[package]
name = "lance-bench"
version = "0.1.0"
edition = "2024"
# Standalone pilot for Phase B (see docs/EXECUTION_PLAN.md).
# Deliberately NOT sharing workspace deps — Lance 4.x pulls in its own
# DataFusion and Arrow versions incompatible with the rest of the stack.
# Isolating the pilot means we don't force a workspace-wide upgrade until
# we've decided Lance is worth it.
[dependencies]
# Only the features we actually need — the default brings in AWS/Azure/GCP/HF etc
# which is ~200 extra crates we don't care about for a local pilot.
lance = { version = "4.0", default-features = false }
# Lance exposes DatasetIndexExt, IndexType, and IvfBuildParams through
# its sub-crates which must be imported directly — lance itself doesn't
# re-export them at a convenient path.
lance-index = { version = "4.0", default-features = false }
lance-linalg = { version = "4.0", default-features = false }
# Arrow re-exported by Lance; pin to a range Lance picks so types match.
arrow = "57"
arrow-array = "57"
arrow-schema = "57"
# Also need to read the EXISTING Parquet vector files so we can compare.
# These live in data/vectors/*.parquet. Lance's internal Parquet reading
# might differ from ours; using our format's Arrow/Parquet versions for
# the read side keeps the inputs identical.
parquet = "57"
tokio = { version = "1", features = ["full"] }
futures = "0.3"
serde = { version = "1", features = ["derive"] }
serde_json = "1"
anyhow = "1"
bytes = "1"
[[bin]]
name = "lance-bench"
path = "src/main.rs"
+633
View File
@@ -0,0 +1,633 @@
//! Phase B: Lance pilot benchmark.
//!
//! Standalone binary that compares Lance vector storage against our
//! Parquet-with-binary-blob + in-RAM HNSW approach. See
//! docs/EXECUTION_PLAN.md for the decision rules this fuels.
//!
//! Inputs:
//! data/vectors/resumes_100k_v2.parquet — existing 100K × 768d embeddings
//!
//! Output:
//! A JSON report printed to stdout with measurements for:
//! - Cold load time (parquet → arrow) vs Lance open + scan
//! - Disk size
//! - Vector search latency (p50 / p95 / p99)
//! - Single-row random access
//! - Append cost (adding 10K rows)
//!
//! Usage:
//! cargo run --bin lance-bench -- \
//! --parquet data/vectors/resumes_100k_v2.parquet \
//! --lance-out /tmp/lance_resumes_100k_v2 \
//! --json-out /tmp/lance_bench.json
use anyhow::{Context, Result};
use arrow_array::{Array, ArrayRef, BinaryArray, FixedSizeListArray, Float32Array, RecordBatch, RecordBatchIterator};
use arrow_schema::{DataType, Field, Schema};
use serde::Serialize;
use std::sync::Arc;
use std::time::Instant;
#[derive(Debug, Serialize)]
struct BenchReport {
vectors: usize,
dimensions: usize,
parquet_path: String,
lance_path: String,
// Parquet baseline
parquet_disk_bytes: u64,
parquet_cold_load_secs: f32,
// Lance numbers
lance_write_secs: f32,
lance_disk_bytes: u64,
lance_cold_open_secs: f32,
// Index + search
lance_index_build_secs: Option<f32>,
lance_index_disk_bytes: Option<u64>,
lance_search_p50_us: Option<f32>,
lance_search_p95_us: Option<f32>,
lance_search_p99_us: Option<f32>,
// Architectural features Parquet+sidecar can't cheaply do
lance_random_row_access_us: Option<f32>, // fetch one row by row_id
parquet_random_row_access_us: Option<f32>, // for comparison — full scan cost
lance_append_10k_secs: Option<f32>, // add 10K new rows
lance_append_disk_bytes_added: Option<u64>,
// Head-to-head reference (from our own measurements)
reference_hnsw_p50_us: f32,
reference_hnsw_p95_us: f32,
reference_brute_force_us: f32,
reference_hnsw_build_secs: f32,
}
#[tokio::main]
async fn main() -> Result<()> {
// Simple positional args: parquet_in, lance_out.
let args: Vec<String> = std::env::args().collect();
let parquet_path = args
.get(1)
.cloned()
.unwrap_or_else(|| "data/vectors/resumes_100k_v2.parquet".to_string());
let lance_path = args
.get(2)
.cloned()
.unwrap_or_else(|| "/tmp/lance_bench_dataset".to_string());
eprintln!("=== Phase B Lance pilot ===");
eprintln!("input parquet: {}", parquet_path);
eprintln!("output lance: {}", lance_path);
// --- 1. Cold-load the existing Parquet vector index into memory
eprintln!("\n[1/4] reading Parquet baseline...");
let t0 = Instant::now();
let (schema, batches, total_rows) = read_parquet_vectors(&parquet_path)
.context("read parquet")?;
let parquet_cold_load_secs = t0.elapsed().as_secs_f32();
let parquet_disk_bytes = std::fs::metadata(&parquet_path)?.len();
let dims = detect_vector_dims(&batches)?;
eprintln!(
" loaded {} rows, {} columns, vectors={}d, disk={:.1} MB, cold load={:.2}s",
total_rows,
schema.fields().len(),
dims,
parquet_disk_bytes as f64 / 1_000_000.0,
parquet_cold_load_secs,
);
// --- 2. Convert from binary-blob-of-f32 to Lance's FixedSizeList<Float32>
eprintln!("\n[2/4] converting binary-blob vectors to Arrow FixedSizeList...");
let t0 = Instant::now();
let (lance_schema, lance_batches) = convert_to_fixed_size_list(&schema, batches, dims)?;
eprintln!(" conversion took {:.2}s", t0.elapsed().as_secs_f32());
// --- 3. Write as Lance dataset
eprintln!("\n[3/4] writing Lance dataset...");
let t0 = Instant::now();
// Clean up any prior run
let _ = std::fs::remove_dir_all(&lance_path);
write_lance_dataset(&lance_path, lance_schema.clone(), lance_batches).await?;
let lance_write_secs = t0.elapsed().as_secs_f32();
let lance_disk_bytes = dir_size_bytes(&lance_path);
eprintln!(
" write took {:.2}s, disk={:.1} MB",
lance_write_secs,
lance_disk_bytes as f64 / 1_000_000.0,
);
// --- 4. Cold open + scan the Lance dataset
eprintln!("\n[4/6] cold-opening Lance dataset...");
let t0 = Instant::now();
let scanned_rows = cold_open_and_scan_lance(&lance_path).await?;
let lance_cold_open_secs = t0.elapsed().as_secs_f32();
eprintln!(
" open + full scan: {} rows in {:.2}s",
scanned_rows, lance_cold_open_secs,
);
// --- 5. Build a vector index on the Lance dataset
eprintln!("\n[5/6] building Lance vector index (IVF_PQ)...");
let t0 = Instant::now();
let index_built = build_lance_vector_index(&lance_path, dims).await;
let (lance_index_build_secs, lance_index_disk_bytes) = match index_built {
Ok(()) => {
let secs = t0.elapsed().as_secs_f32();
let disk = dir_size_bytes(&lance_path) - lance_disk_bytes;
eprintln!(" built in {:.2}s, index adds {:.1} MB on disk", secs, disk as f64 / 1e6);
(Some(secs), Some(disk))
}
Err(e) => {
eprintln!(" index build failed: {e:#}");
(None, None)
}
};
// --- 6. Run search queries, measure latency
eprintln!("\n[6/6] running vector search benchmarks...");
let search_stats = if lance_index_build_secs.is_some() {
run_search_benchmarks(&lance_path, dims).await.ok()
} else {
None
};
let (lance_search_p50, lance_search_p95, lance_search_p99) = match search_stats {
Some((p50, p95, p99)) => {
eprintln!(" p50={:.0}us p95={:.0}us p99={:.0}us", p50, p95, p99);
(Some(p50), Some(p95), Some(p99))
}
None => (None, None, None),
};
// --- Random access comparison
eprintln!("\n[7/8] random row access — Lance vs full-scan Parquet...");
let lance_random = measure_random_access_lance(&lance_path).await.ok();
let parquet_random = measure_random_access_parquet(&parquet_path).ok();
if let Some(us) = lance_random {
eprintln!(" Lance random-fetch avg: {:.0}us", us);
}
if let Some(us) = parquet_random {
eprintln!(" Parquet full-scan-to-row avg: {:.0}us", us);
}
// --- Append cost
eprintln!("\n[8/8] append 10K new rows to existing dataset...");
let t0 = Instant::now();
let pre_append_bytes = dir_size_bytes(&lance_path);
let append_result = append_10k_rows(&lance_path, dims).await;
let (lance_append_secs, lance_append_bytes) = match append_result {
Ok(()) => {
let secs = t0.elapsed().as_secs_f32();
let bytes = dir_size_bytes(&lance_path).saturating_sub(pre_append_bytes);
eprintln!(" append took {:.2}s, added {:.1} MB", secs, bytes as f64 / 1e6);
(Some(secs), Some(bytes))
}
Err(e) => {
eprintln!(" append failed: {e:#}");
(None, None)
}
};
// --- Report
let report = BenchReport {
vectors: total_rows,
dimensions: dims,
parquet_path: parquet_path.clone(),
lance_path: lance_path.clone(),
parquet_disk_bytes,
parquet_cold_load_secs,
lance_write_secs,
lance_disk_bytes,
lance_cold_open_secs,
lance_index_build_secs,
lance_index_disk_bytes,
lance_search_p50_us: lance_search_p50,
lance_search_p95_us: lance_search_p95,
lance_search_p99_us: lance_search_p99,
lance_random_row_access_us: lance_random,
parquet_random_row_access_us: parquet_random,
lance_append_10k_secs: lance_append_secs,
lance_append_disk_bytes_added: lance_append_bytes,
// From our Phase 15 trial on the SAME index (ec=80 es=30, recall=1.00):
reference_hnsw_p50_us: 873.0,
reference_hnsw_p95_us: 1413.0,
reference_brute_force_us: 43983.0,
reference_hnsw_build_secs: 230.0,
};
let json = serde_json::to_string_pretty(&report)?;
println!("{}", json);
eprintln!("\n=== Summary ===");
eprintln!(" Parquet cold load: {:.2}s", report.parquet_cold_load_secs);
eprintln!(" Lance cold open: {:.2}s ({})",
report.lance_cold_open_secs,
format_ratio(report.parquet_cold_load_secs, report.lance_cold_open_secs));
eprintln!(" Parquet disk: {:.1} MB", report.parquet_disk_bytes as f64 / 1e6);
eprintln!(" Lance disk: {:.1} MB ({})",
report.lance_disk_bytes as f64 / 1e6,
format_ratio(report.parquet_disk_bytes as f32, report.lance_disk_bytes as f32));
if let (Some(p50), Some(p95)) = (report.lance_search_p50_us, report.lance_search_p95_us) {
eprintln!(" Lance search p50: {:.0}us vs our HNSW {:.0}us ({})",
p50, report.reference_hnsw_p50_us,
format_ratio(report.reference_hnsw_p50_us, p50));
eprintln!(" Lance search p95: {:.0}us vs our HNSW {:.0}us ({})",
p95, report.reference_hnsw_p95_us,
format_ratio(report.reference_hnsw_p95_us, p95));
eprintln!(" Speedup vs brute force: {:.1}× (Lance) vs {:.1}× (HNSW)",
report.reference_brute_force_us / p50,
report.reference_brute_force_us / report.reference_hnsw_p50_us);
}
if let Some(build) = report.lance_index_build_secs {
eprintln!(" Index build: {:.1}s (Lance IVF_PQ) vs {:.0}s (our HNSW ec=80) ({}× faster)",
build, report.reference_hnsw_build_secs, report.reference_hnsw_build_secs / build);
}
if let (Some(lance_us), Some(parquet_us)) = (report.lance_random_row_access_us, report.parquet_random_row_access_us) {
eprintln!(" Random row access: {:.0}us (Lance) vs {:.0}us (Parquet scan) ({})",
lance_us, parquet_us, format_ratio(parquet_us, lance_us));
}
if let Some(append_secs) = report.lance_append_10k_secs {
eprintln!(" Append 10K rows: {:.2}s (Lance native) [Parquet would require full rewrite]",
append_secs);
}
Ok(())
}
fn format_ratio(baseline: f32, candidate: f32) -> String {
if candidate == 0.0 { return "inf".into(); }
let ratio = baseline / candidate;
if ratio >= 1.0 {
format!("{:.2}× faster/smaller", ratio)
} else {
format!("{:.2}× slower/larger", 1.0 / ratio)
}
}
fn dir_size_bytes(path: &str) -> u64 {
fn recurse(p: &std::path::Path) -> u64 {
let Ok(meta) = std::fs::metadata(p) else { return 0; };
if meta.is_file() { return meta.len(); }
let Ok(entries) = std::fs::read_dir(p) else { return 0; };
entries
.filter_map(|e| e.ok())
.map(|e| recurse(&e.path()))
.sum()
}
recurse(std::path::Path::new(path))
}
/// Read the existing vector Parquet (binary-blob format: source, doc_id,
/// chunk_idx, chunk_text, vector as Binary bytes).
fn read_parquet_vectors(path: &str) -> Result<(Arc<Schema>, Vec<RecordBatch>, usize)> {
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use std::fs::File;
let file = File::open(path).with_context(|| format!("open {path}"))?;
let builder = ParquetRecordBatchReaderBuilder::try_new(file)?;
let schema = builder.schema().clone();
let reader = builder.build()?;
let batches: Vec<RecordBatch> = reader.collect::<Result<Vec<_>, _>>()?;
let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
Ok((schema, batches, rows))
}
fn detect_vector_dims(batches: &[RecordBatch]) -> Result<usize> {
for batch in batches {
let vector_col_idx = batch
.schema()
.index_of("vector")
.context("no 'vector' column in parquet")?;
let col = batch.column(vector_col_idx);
if let Some(binary) = col.as_any().downcast_ref::<BinaryArray>() {
for i in 0..binary.len() {
if !binary.is_null(i) {
let bytes = binary.value(i);
return Ok(bytes.len() / 4); // f32 = 4 bytes
}
}
}
}
anyhow::bail!("could not determine vector dimensions")
}
/// Convert our binary-blob vector representation into Arrow's native
/// FixedSizeList<Float32> — that's what Lance expects for vector columns.
fn convert_to_fixed_size_list(
schema: &Arc<Schema>,
batches: Vec<RecordBatch>,
dims: usize,
) -> Result<(Arc<Schema>, Vec<RecordBatch>)> {
// New schema keeps everything identical but replaces the vector column
// with a FixedSizeList<Float32, dims>.
let new_fields: Vec<Arc<Field>> = schema
.fields()
.iter()
.map(|f| {
if f.name() == "vector" {
Arc::new(Field::new(
"vector",
DataType::FixedSizeList(
Arc::new(Field::new("item", DataType::Float32, true)),
dims as i32,
),
false,
))
} else {
f.clone()
}
})
.collect();
let new_schema = Arc::new(Schema::new(new_fields));
let mut new_batches = Vec::with_capacity(batches.len());
for batch in batches {
let vector_idx = batch.schema().index_of("vector")?;
let mut new_arrays: Vec<ArrayRef> = Vec::with_capacity(batch.num_columns());
for (i, col) in batch.columns().iter().enumerate() {
if i == vector_idx {
let binary = col
.as_any()
.downcast_ref::<BinaryArray>()
.context("vector column must be Binary")?;
let fsl = binary_to_fixed_size_list(binary, dims)?;
new_arrays.push(Arc::new(fsl));
} else {
new_arrays.push(col.clone());
}
}
new_batches.push(RecordBatch::try_new(new_schema.clone(), new_arrays)?);
}
Ok((new_schema, new_batches))
}
fn binary_to_fixed_size_list(binary: &BinaryArray, dims: usize) -> Result<FixedSizeListArray> {
let n = binary.len();
let mut all_floats: Vec<f32> = Vec::with_capacity(n * dims);
for i in 0..n {
if binary.is_null(i) {
all_floats.extend(std::iter::repeat(0.0).take(dims));
continue;
}
let bytes = binary.value(i);
if bytes.len() != dims * 4 {
anyhow::bail!(
"row {} has {} bytes, expected {} ({} × f32)",
i, bytes.len(), dims * 4, dims,
);
}
for chunk in bytes.chunks_exact(4) {
all_floats.push(f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]));
}
}
let values = Float32Array::from(all_floats);
let field = Arc::new(Field::new("item", DataType::Float32, true));
FixedSizeListArray::try_new(field, dims as i32, Arc::new(values), None)
.context("build FixedSizeListArray")
}
/// Write batches into a Lance dataset at the given path.
async fn write_lance_dataset(
path: &str,
schema: Arc<Schema>,
batches: Vec<RecordBatch>,
) -> Result<()> {
use lance::dataset::{Dataset, WriteParams};
let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema);
Dataset::write(reader, path, Some(WriteParams::default()))
.await
.context("Dataset::write")?;
Ok(())
}
/// Open a Lance dataset cold (from disk) and scan it fully — measuring the
/// equivalent of our "load embeddings from Parquet" cost.
async fn cold_open_and_scan_lance(path: &str) -> Result<usize> {
use futures::StreamExt;
use lance::dataset::Dataset;
let dataset = Dataset::open(path).await.context("Dataset::open")?;
let scanner = dataset.scan();
let mut stream = scanner.try_into_stream().await?;
let mut total = 0usize;
while let Some(batch) = stream.next().await {
let batch = batch?;
total += batch.num_rows();
}
Ok(total)
}
/// Build an IVF_PQ vector index on the `vector` column. IVF_PQ (Inverted File
/// with Product Quantization) is Lance's native ANN index — comparable to
/// HNSW in intent, but on-disk and compatible with Lance's random-access
/// model.
async fn build_lance_vector_index(path: &str, _dims: usize) -> Result<()> {
use lance::dataset::Dataset;
use lance::index::vector::VectorIndexParams;
use lance_index::{DatasetIndexExt, IndexType};
use lance_linalg::distance::MetricType;
let mut dataset = Dataset::open(path).await?;
// IVF_PQ with ~sqrt(N) partitions is a reasonable default for 100K.
// num_sub_vectors must divide dims evenly: 768/48 = 16 dims per subvector.
// num_bits = 8 gives 256 codes per subvector (good recall/size trade).
// max_iterations = 50 is plenty for this scale.
let params = VectorIndexParams::ivf_pq(
316, // num_partitions (~sqrt(100000))
8, // num_bits
48, // num_sub_vectors
MetricType::Cosine,
50, // max_iterations
);
dataset
.create_index(
&["vector"],
IndexType::Vector,
Some("vec_idx".into()),
&params,
true,
)
.await
.context("create_index")?;
Ok(())
}
/// Run N vector searches against the Lance dataset and return (p50, p95, p99) latencies in us.
/// Uses a handful of random rows as queries — same pattern as our harness::synthetic_from_chunks.
async fn run_search_benchmarks(path: &str, _dims: usize) -> Result<(f32, f32, f32)> {
use futures::StreamExt;
use lance::dataset::Dataset;
let dataset = Dataset::open(path).await?;
// Pick 20 representative query vectors from the data itself.
// (Synthetic — same pattern as our existing harness.)
let query_vectors = sample_query_vectors(&dataset, 20).await?;
let mut latencies_us: Vec<f32> = Vec::with_capacity(query_vectors.len());
for (i, qv) in query_vectors.iter().enumerate() {
let qarr = Arc::new(Float32Array::from(qv.clone())) as ArrayRef;
let t0 = Instant::now();
let mut scanner = dataset.scan();
scanner
.nearest("vector", qarr.as_any().downcast_ref::<Float32Array>().unwrap(), 10)
.context("scanner.nearest")?;
let mut stream = scanner.try_into_stream().await?;
let mut hits = 0;
while let Some(batch) = stream.next().await {
let batch = batch?;
hits += batch.num_rows();
}
let us = t0.elapsed().as_micros() as f32;
latencies_us.push(us);
if i == 0 {
eprintln!(" first query: {} hits in {:.0}us (includes any lazy init)", hits, us);
}
}
latencies_us.sort_by(|a, b| a.partial_cmp(b).unwrap());
let p = |pct: f32| -> f32 {
let idx = ((latencies_us.len() as f32 - 1.0) * pct).round() as usize;
latencies_us[idx.min(latencies_us.len() - 1)]
};
Ok((p(0.50), p(0.95), p(0.99)))
}
/// Random row access via Lance's `take` — fetch 20 random rows by index, measure avg latency.
async fn measure_random_access_lance(path: &str) -> Result<f32> {
use lance::dataset::Dataset;
let dataset = Dataset::open(path).await?;
let n = dataset.count_rows(None).await?;
let indices: Vec<u64> = (0..20).map(|i| ((i as u64) * (n as u64 / 23)) % (n as u64)).collect();
// Full-schema projection — Lance's Schema implements Into<ProjectionRequest>.
let schema = dataset.schema().clone();
let mut total_us: u128 = 0;
for idx in &indices {
let t0 = Instant::now();
let _batch = dataset.take(&[*idx], schema.clone()).await?;
total_us += t0.elapsed().as_micros();
}
Ok(total_us as f32 / indices.len() as f32)
}
/// Random row access for Parquet — full scan + filter. There's no random-access
/// primitive in vanilla Parquet, so this is the cost of finding one specific row.
/// This is the cost our current design pays for "get doc X's full text for RAG."
fn measure_random_access_parquet(path: &str) -> Result<f32> {
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use std::fs::File;
// We simulate 5 lookups — full scan each time. 20 would be painful.
let iters = 5;
let mut total_us: u128 = 0;
for _ in 0..iters {
let t0 = Instant::now();
let file = File::open(path)?;
let builder = ParquetRecordBatchReaderBuilder::try_new(file)?;
let reader = builder.build()?;
// Iterate until we've conceptually found a row — we stop early if
// we wanted row 50000, but we have to at least read its batch.
let mut seen = 0usize;
for b in reader {
let b = b?;
seen += b.num_rows();
if seen > 50000 { break; }
}
total_us += t0.elapsed().as_micros();
}
Ok(total_us as f32 / iters as f32)
}
/// Append 10K new rows to the existing Lance dataset.
/// Measures the "ingest delta" cost without full rewrite.
async fn append_10k_rows(path: &str, dims: usize) -> Result<()> {
use lance::dataset::{Dataset, WriteMode, WriteParams};
let dataset = Dataset::open(path).await?;
let schema = dataset.schema();
let arrow_schema: Arc<Schema> = Arc::new(schema.into());
// Build a 10K row batch with random-ish data matching the existing schema.
let n = 10_000;
let arrays: Vec<ArrayRef> = arrow_schema
.fields()
.iter()
.map(|f| -> Result<ArrayRef> {
match f.data_type() {
DataType::Utf8 => {
let vals: Vec<String> = (0..n).map(|i| format!("appended-{}", i)).collect();
Ok(Arc::new(arrow_array::StringArray::from(vals)))
}
DataType::Int32 => {
let vals: Vec<i32> = (0..n as i32).collect();
Ok(Arc::new(arrow_array::Int32Array::from(vals)))
}
DataType::FixedSizeList(_, _) => {
let floats: Vec<f32> = (0..n * dims).map(|i| (i as f32).sin()).collect();
let values = Float32Array::from(floats);
let field = Arc::new(Field::new("item", DataType::Float32, true));
let fsl = FixedSizeListArray::try_new(field, dims as i32, Arc::new(values), None)?;
Ok(Arc::new(fsl))
}
other => anyhow::bail!("unsupported append column type: {:?}", other),
}
})
.collect::<Result<Vec<_>>>()?;
let batch = RecordBatch::try_new(arrow_schema.clone(), arrays)?;
let reader = RecordBatchIterator::new(vec![Ok(batch)].into_iter(), arrow_schema);
let params = WriteParams { mode: WriteMode::Append, ..Default::default() };
Dataset::write(reader, path, Some(params)).await?;
Ok(())
}
/// Grab a few existing vectors from the dataset to use as self-similar queries.
async fn sample_query_vectors(
dataset: &lance::dataset::Dataset,
count: usize,
) -> Result<Vec<Vec<f32>>> {
use futures::StreamExt;
// Just take the first `count` rows; good enough for latency measurement.
let scanner = dataset.scan();
let mut scanner = scanner;
scanner.limit(Some(count as i64), None)?;
scanner.project(&["vector"])?;
let mut stream = scanner.try_into_stream().await?;
let mut out = Vec::with_capacity(count);
while let Some(batch) = stream.next().await {
let batch = batch?;
let vector_col = batch
.column(0)
.as_any()
.downcast_ref::<FixedSizeListArray>()
.context("vector column must be FixedSizeList")?;
for row in 0..vector_col.len() {
if out.len() >= count { break; }
let values = vector_col.value(row);
let f32_arr = values
.as_any()
.downcast_ref::<Float32Array>()
.context("inner array must be Float32")?;
let mut v = Vec::with_capacity(f32_arr.len());
for i in 0..f32_arr.len() {
v.push(f32_arr.value(i));
}
out.push(v);
}
if out.len() >= count { break; }
}
Ok(out)
}
+11
View File
@@ -0,0 +1,11 @@
[package]
name = "proto"
version = "0.1.0"
edition = "2024"
[dependencies]
tonic = { workspace = true }
prost = { workspace = true }
[build-dependencies]
tonic-build = { workspace = true }
+4
View File
@@ -0,0 +1,4 @@
fn main() -> Result<(), Box<dyn std::error::Error>> {
tonic_build::compile_protos("../../proto/lakehouse.proto")?;
Ok(())
}
+3
View File
@@ -0,0 +1,3 @@
pub mod lakehouse {
tonic::include_proto!("lakehouse");
}
+23
View File
@@ -0,0 +1,23 @@
[package]
name = "queryd"
version = "0.1.0"
edition = "2024"
[dependencies]
shared = { path = "../shared" }
catalogd = { path = "../catalogd" }
storaged = { path = "../storaged" }
truth = { path = "../truth" }
tokio = { workspace = true }
axum = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }
tracing = { workspace = true }
datafusion = { workspace = true }
object_store = { workspace = true }
arrow = { workspace = true }
parquet = { workspace = true }
bytes = { workspace = true }
futures = { workspace = true }
url = { workspace = true }
chrono = { workspace = true }
+135
View File
@@ -0,0 +1,135 @@
/// In-memory cache for hot datasets.
/// Pinned datasets are loaded as Arrow RecordBatches and served from RAM.
/// LRU eviction keeps memory bounded.
use arrow::array::RecordBatch;
use arrow::datatypes::SchemaRef;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
/// A cached dataset — schema + all batches in memory.
#[derive(Clone)]
pub struct CachedDataset {
pub name: String,
pub schema: SchemaRef,
pub batches: Vec<RecordBatch>,
pub row_count: usize,
pub size_bytes: usize, // approximate
pub last_accessed: std::time::Instant,
}
/// Memory cache with LRU eviction.
#[derive(Clone)]
pub struct MemCache {
datasets: Arc<RwLock<HashMap<String, CachedDataset>>>,
max_bytes: usize,
}
impl MemCache {
pub fn new(max_bytes: usize) -> Self {
Self {
datasets: Arc::new(RwLock::new(HashMap::new())),
max_bytes,
}
}
/// Load a dataset into the cache.
pub async fn put(&self, name: &str, schema: SchemaRef, batches: Vec<RecordBatch>) {
let row_count: usize = batches.iter().map(|b| b.num_rows()).sum();
let size_bytes: usize = batches.iter()
.map(|b| b.get_array_memory_size())
.sum();
// Evict if needed
self.evict_to_fit(size_bytes).await;
let entry = CachedDataset {
name: name.to_string(),
schema,
batches,
row_count,
size_bytes,
last_accessed: std::time::Instant::now(),
};
tracing::info!("cached '{}': {} rows, {:.1} MB",
name, row_count, size_bytes as f64 / 1024.0 / 1024.0);
let mut cache = self.datasets.write().await;
cache.insert(name.to_string(), entry);
}
/// Get a cached dataset, updating access time.
pub async fn get(&self, name: &str) -> Option<CachedDataset> {
let mut cache = self.datasets.write().await;
if let Some(entry) = cache.get_mut(name) {
entry.last_accessed = std::time::Instant::now();
Some(entry.clone())
} else {
None
}
}
/// Check if a dataset is cached.
pub async fn contains(&self, name: &str) -> bool {
self.datasets.read().await.contains_key(name)
}
/// Evict a specific dataset.
pub async fn evict(&self, name: &str) -> bool {
let mut cache = self.datasets.write().await;
cache.remove(name).is_some()
}
/// Get cache stats.
pub async fn stats(&self) -> CacheStats {
let cache = self.datasets.read().await;
let total_bytes: usize = cache.values().map(|d| d.size_bytes).sum();
let total_rows: usize = cache.values().map(|d| d.row_count).sum();
CacheStats {
datasets: cache.len(),
total_bytes,
total_rows,
max_bytes: self.max_bytes,
names: cache.keys().cloned().collect(),
}
}
/// Evict least-recently-used entries to make room for `needed_bytes`.
async fn evict_to_fit(&self, needed_bytes: usize) {
let mut cache = self.datasets.write().await;
let current: usize = cache.values().map(|d| d.size_bytes).sum();
if current + needed_bytes <= self.max_bytes {
return;
}
// Sort by last accessed (oldest first)
let mut entries: Vec<(String, std::time::Instant, usize)> = cache.iter()
.map(|(k, v)| (k.clone(), v.last_accessed, v.size_bytes))
.collect();
entries.sort_by_key(|(_, t, _)| *t);
let mut freed = 0usize;
let target = (current + needed_bytes).saturating_sub(self.max_bytes);
for (name, _, size) in entries {
if freed >= target {
break;
}
tracing::info!("evicting '{}' from cache ({:.1} MB)", name, size as f64 / 1024.0 / 1024.0);
cache.remove(&name);
freed += size;
}
}
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct CacheStats {
pub datasets: usize,
pub total_bytes: usize,
pub total_rows: usize,
pub max_bytes: usize,
pub names: Vec<String>,
}
+328
View File
@@ -0,0 +1,328 @@
use catalogd::registry::Registry;
use datafusion::datasource::listing::{ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl};
use datafusion::datasource::MemTable;
use datafusion::datasource::file_format::parquet::ParquetFormat;
use datafusion::prelude::*;
use object_store::ObjectStore;
use std::sync::Arc;
use storaged::registry::BucketRegistry;
use url::Url;
use crate::cache::MemCache;
use crate::delta;
const STORE_SCHEME: &str = "lakehouse";
/// URL scheme used to register a bucket's object store with DataFusion.
/// Bucket names safely embed in URL host components after sanitizing
/// `:` (used in `profile:user`) to `-`.
fn bucket_scheme(bucket: &str) -> String {
format!("{STORE_SCHEME}-{}", bucket.replace(':', "-"))
}
/// Build a `ListingTableUrl`-shaped string for an object in a bucket.
fn bucket_object_url(bucket: &str, key: &str) -> String {
format!("{}://data/{}", bucket_scheme(bucket), key)
}
/// Build the registration base URL for a bucket.
fn bucket_base_url(bucket: &str) -> String {
format!("{}://data/", bucket_scheme(bucket))
}
/// Query engine with in-memory cache and delta merge support.
/// Federation layer 2: holds a `BucketRegistry` so cross-bucket queries
/// route reads to the right ObjectStore based on `ObjectRef.bucket`.
#[derive(Clone)]
pub struct QueryEngine {
registry: Registry,
buckets: Arc<BucketRegistry>,
/// Primary bucket store cached for backward-compat callers (compact,
/// workspace ops) that haven't been migrated to bucket-aware paths.
store: Arc<dyn ObjectStore>,
cache: MemCache,
}
impl QueryEngine {
pub fn new(registry: Registry, buckets: Arc<BucketRegistry>, cache: MemCache) -> Self {
let store = buckets.default_store();
Self { registry, buckets, store, cache }
}
pub fn cache(&self) -> &MemCache {
&self.cache
}
/// Backward-compat: returns the primary bucket's store. Bucket-aware
/// callers should use `engine.buckets()` directly.
pub fn store(&self) -> &Arc<dyn ObjectStore> {
&self.store
}
pub fn buckets(&self) -> &Arc<BucketRegistry> {
&self.buckets
}
/// Catalog registry — used by compact/tombstone integration.
pub fn catalog(&self) -> &Registry {
&self.registry
}
/// Execute a SQL query. Uses cache for hot data, falls back to Parquet.
pub async fn query(&self, sql: &str) -> Result<Vec<arrow::array::RecordBatch>, String> {
let ctx = self.build_context().await?;
let df = ctx.sql(sql).await.map_err(|e| format!("SQL error: {e}"))?;
let batches = df.collect().await.map_err(|e| format!("execution error: {e}"))?;
Ok(batches)
}
/// Pin a dataset into the memory cache.
pub async fn pin_dataset(&self, name: &str) -> Result<(), String> {
let ctx = SessionContext::new();
// Register every bucket so a multi-bucket dataset can be pinned.
for info in self.buckets.list().await {
let url = Url::parse(&bucket_base_url(&info.name))
.map_err(|e| format!("invalid store url: {e}"))?;
let store = self.buckets.get(&info.name)
.map_err(|e| format!("registry inconsistency: {e}"))?;
ctx.runtime_env().register_object_store(&url, store);
}
let dataset = self.registry.get_by_name(name).await
.ok_or_else(|| format!("dataset not found: {name}"))?;
if dataset.objects.is_empty() {
return Err(format!("dataset '{name}' has no objects"));
}
let opts = ListingOptions::new(Arc::new(ParquetFormat::default()));
let table_paths: Vec<ListingTableUrl> = dataset.objects.iter()
.filter_map(|o| {
let bucket = if o.bucket.is_empty() { "primary" } else { &o.bucket };
ListingTableUrl::parse(&bucket_object_url(bucket, &o.key)).ok()
})
.collect();
let schema = opts.infer_schema(&ctx.state(), &table_paths[0]).await
.map_err(|e| format!("schema inference: {e}"))?;
let config = ListingTableConfig::new_with_multi_paths(table_paths)
.with_listing_options(opts)
.with_schema(schema.clone());
let table = ListingTable::try_new(config).map_err(|e| format!("table: {e}"))?;
ctx.register_table(name, Arc::new(table)).map_err(|e| format!("register: {e}"))?;
let df = ctx.sql(&format!("SELECT * FROM {name}")).await.map_err(|e| format!("read: {e}"))?;
let batches = df.collect().await.map_err(|e| format!("collect: {e}"))?;
self.cache.put(name, schema, batches).await;
Ok(())
}
async fn build_context(&self) -> Result<SessionContext, String> {
let ctx = SessionContext::new();
// Phase E: snapshot tombstones by dataset before registering tables
// so we can wrap tombstoned tables in a filter view. The underlying
// base table is registered under an internal name `__raw__{dataset}`
// and the public `{dataset}` name becomes the filtered view.
let all_dataset_names: Vec<String> = self.registry.list().await
.iter().map(|d| d.name.clone()).collect();
let tombstones_by_dataset = self
.registry
.tombstones()
.all_grouped(&all_dataset_names)
.await
.unwrap_or_default();
// Federation layer 2: register every configured bucket as its own
// DataFusion ObjectStore under a distinct URL scheme. Each
// dataset's ObjectRef.bucket determines which store DataFusion
// routes to at scan time.
let bucket_infos = self.buckets.list().await;
for info in &bucket_infos {
let url_str = bucket_base_url(&info.name);
let url = Url::parse(&url_str)
.map_err(|e| format!("invalid store url '{url_str}': {e}"))?;
// unknown-bucket here would be a config invariant violation,
// since we just listed them — propagate as an error to surface it.
let store = self.buckets.get(&info.name)
.map_err(|e| format!("registry inconsistency: {e}"))?;
ctx.runtime_env().register_object_store(&url, store);
}
// Backward-compat: also register the legacy `lakehouse://data/`
// URL pointing at primary, in case any code path still constructs
// ListingTableUrls without a bucket prefix.
let legacy_url = Url::parse(&format!("{STORE_SCHEME}://data/"))
.map_err(|e| format!("invalid legacy store url: {e}"))?;
ctx.runtime_env().register_object_store(&legacy_url, self.store.clone());
let datasets = self.registry.list().await;
for dataset in &datasets {
if dataset.objects.is_empty() {
continue;
}
// Check cache first — cached batches are bucket-agnostic since
// they're already materialized in memory.
if let Some(cached) = self.cache.get(&dataset.name).await {
let delta_batches = delta::load_deltas(&self.store, &dataset.name).await.unwrap_or_default();
let mut all_batches = cached.batches;
all_batches.extend(delta_batches);
let mem_table = MemTable::try_new(cached.schema, vec![all_batches])
.map_err(|e| format!("MemTable error for {}: {e}", dataset.name))?;
ctx.register_table(&dataset.name, Arc::new(mem_table))
.map_err(|e| format!("register cached {}: {e}", dataset.name))?;
tracing::debug!("using cached table: {}", dataset.name);
continue;
}
// Build ListingTable URLs that include the per-object bucket
// scheme. DataFusion routes each scan to the right store
// automatically based on the URL.
let opts = ListingOptions::new(Arc::new(ParquetFormat::default()));
let table_paths: Vec<ListingTableUrl> = dataset.objects.iter()
.filter_map(|o| {
let bucket = if o.bucket.is_empty() { "primary" } else { &o.bucket };
ListingTableUrl::parse(&bucket_object_url(bucket, &o.key)).ok()
})
.collect();
if table_paths.is_empty() {
continue;
}
let schema = opts.infer_schema(&ctx.state(), &table_paths[0]).await
.map_err(|e| format!("schema inference failed for {}: {e}", dataset.name))?;
let config = ListingTableConfig::new_with_multi_paths(table_paths)
.with_listing_options(opts)
.with_schema(schema);
let table = ListingTable::try_new(config)
.map_err(|e| format!("table creation failed for {}: {e}", dataset.name))?;
// Decide the registration name: if this dataset has any
// tombstones, the raw table gets an internal name and the
// public name becomes a filtered view.
let tombstone_entry = tombstones_by_dataset.get(&dataset.name);
let register_name = if tombstone_entry.is_some() {
format!("__raw__{}", dataset.name)
} else {
dataset.name.clone()
};
if let Err(e) = ctx.register_table(register_name.as_str(), Arc::new(table)) {
let msg = e.to_string();
if msg.contains("already exists") {
tracing::debug!("skip duplicate manifest registration: {}", register_name);
continue;
} else {
return Err(format!("table registration failed for {}: {}", register_name, msg));
}
}
// If there are tombstones, register the public name as a
// filtered view that excludes tombstoned row_key_values.
if let Some((key_col, values)) = tombstone_entry {
// Build WHERE NOT IN (...) — quote values to be SQL-safe.
// For string keys this is a literal list; for integer keys
// CAST(col AS VARCHAR) makes the comparison unambiguous.
let quoted: Vec<String> = values.iter()
.map(|v| format!("'{}'", v.replace('\'', "''")))
.collect();
let sql = format!(
"SELECT * FROM \"{}\" WHERE CAST(\"{}\" AS VARCHAR) NOT IN ({})",
register_name, key_col, quoted.join(", "),
);
tracing::debug!(
"tombstone filter for '{}': {} row_keys excluded",
dataset.name, values.len(),
);
match ctx.sql(&sql).await {
Ok(df) => {
if let Err(e) = ctx.register_table(dataset.name.as_str(), df.into_view()) {
let msg = e.to_string();
if msg.contains("already exists") {
tracing::debug!("skip duplicate tombstone view: {}", dataset.name);
} else {
tracing::warn!("tombstone view registration failed for {}: {msg}", dataset.name);
}
}
}
Err(e) => {
tracing::warn!("tombstone view SQL failed for '{}': {e}", dataset.name);
}
}
}
}
// Phase D: register every AiView as a DataFusion view that wraps
// the base table with the safe projection + filter + redactions.
for view in self.registry.list_views().await {
// Build the SELECT clause with column whitelist and any
// redaction expressions. Each column is either the raw column
// name or an expression aliased back to the column name so
// downstream queries see the same shape.
let select_cols: Vec<String> = view.columns.iter().map(|col| {
match view.column_redactions.get(col) {
None => format!("\"{}\"", col),
Some(shared::types::Redaction::Null) => {
format!("CAST(NULL AS VARCHAR) AS \"{}\"", col)
}
Some(shared::types::Redaction::Hash) => {
// DataFusion has digest('sha256', value) — but column
// type might not be string. Cast first to be safe.
format!("digest(CAST(\"{}\" AS VARCHAR), 'sha256') AS \"{}\"", col, col)
}
Some(shared::types::Redaction::Mask { keep_prefix, keep_suffix }) => {
// Keep first N + last M chars, replace middle with stars.
// CASE NULL through with COALESCE for null-safe behavior.
format!(
"CASE WHEN \"{c}\" IS NULL THEN NULL ELSE \
concat(\
substr(CAST(\"{c}\" AS VARCHAR), 1, {p}), \
repeat('*', greatest(0, length(CAST(\"{c}\" AS VARCHAR)) - {p} - {s})), \
substr(CAST(\"{c}\" AS VARCHAR), greatest(1, length(CAST(\"{c}\" AS VARCHAR)) - {s} + 1))\
) END AS \"{c}\"",
c = col, p = keep_prefix, s = keep_suffix,
)
}
}
}).collect();
let where_clause = view.row_filter
.as_deref()
.map(|f| format!(" WHERE {}", f))
.unwrap_or_default();
let view_sql = format!(
"SELECT {} FROM \"{}\"{}",
select_cols.join(", "),
view.base_dataset,
where_clause,
);
tracing::debug!("registering AiView '{}': {}", view.name, view_sql);
let df = match ctx.sql(&view_sql).await {
Ok(df) => df,
Err(e) => {
tracing::warn!("view '{}' SQL invalid, skipping: {e}", view.name);
continue;
}
};
if let Err(e) = ctx.register_table(view.name.as_str(), df.into_view()) {
let msg = e.to_string();
if msg.contains("already exists") {
tracing::debug!("skip duplicate view registration: {}", view.name);
} else {
tracing::warn!("view '{}' registration failed: {e}", view.name);
}
}
}
Ok(ctx)
}
}
+272
View File
@@ -0,0 +1,272 @@
/// Delta store for incremental updates.
/// Instead of rewriting an entire Parquet file to change one row,
/// we write small delta files. At query time, deltas are merged with the base.
/// Periodic compaction merges deltas into the base file.
use arrow::array::{Array, BooleanArray, Int32Array, Int64Array, RecordBatch, StringArray};
use arrow::compute::filter_record_batch;
use bytes::Bytes;
use object_store::ObjectStore;
use parquet::arrow::ArrowWriter;
use parquet::basic::Compression;
use parquet::file::properties::WriterProperties;
use shared::types::Tombstone;
use std::collections::HashSet;
use std::sync::Arc;
use shared::arrow_helpers::{parquet_to_record_batches, record_batch_to_parquet};
use storaged::ops;
/// Write a delta file for a dataset (new/updated rows).
pub async fn write_delta(
store: &Arc<dyn ObjectStore>,
dataset_name: &str,
batch: &RecordBatch,
) -> Result<String, String> {
let ts = chrono::Utc::now().timestamp_millis();
let key = format!("deltas/{dataset_name}/{ts}.parquet");
let parquet = record_batch_to_parquet(batch)?;
ops::put(store, &key, parquet).await?;
tracing::info!("wrote delta for '{}': {} rows at {}", dataset_name, batch.num_rows(), key);
Ok(key)
}
/// List all delta files for a dataset.
pub async fn list_deltas(
store: &Arc<dyn ObjectStore>,
dataset_name: &str,
) -> Result<Vec<String>, String> {
let prefix = format!("deltas/{dataset_name}/");
ops::list(store, Some(&prefix)).await
}
/// Load all delta batches for a dataset.
pub async fn load_deltas(
store: &Arc<dyn ObjectStore>,
dataset_name: &str,
) -> Result<Vec<RecordBatch>, String> {
let keys = list_deltas(store, dataset_name).await?;
let mut all_batches = Vec::new();
for key in &keys {
let data = ops::get(store, key).await?;
let (_, batches) = parquet_to_record_batches(&data)?;
all_batches.extend(batches);
}
if !all_batches.is_empty() {
let total_rows: usize = all_batches.iter().map(|b| b.num_rows()).sum();
tracing::debug!("loaded {} delta files ({} rows) for '{}'", keys.len(), total_rows, dataset_name);
}
Ok(all_batches)
}
/// Compact: merge base Parquet + all deltas into a single new base file.
/// Optionally filters out tombstoned rows — this is the physical-deletion
/// half of Phase E (the query-time filter handles the "immediate hide",
/// compaction handles the "actually remove from disk" requirement).
///
/// Tombstones must share one `row_key_column` (enforced at tombstone write).
/// If non-empty, every row in the merged output whose row_key_column value
/// matches a tombstone is dropped before the new base is written.
pub async fn compact(
store: &Arc<dyn ObjectStore>,
dataset_name: &str,
base_key: &str,
primary_key_col: Option<&str>,
tombstones: &[Tombstone],
) -> Result<CompactResult, String> {
// Load base
let base_data = ops::get(store, base_key).await?;
let (_schema, mut base_batches) = parquet_to_record_batches(&base_data)?;
// Load deltas
let delta_batches = load_deltas(store, dataset_name).await?;
let delta_count = delta_batches.len();
// Row counts captured before extend; previously base_rows subtracted delta_count (files) from rows — unit mismatch.
let base_row_count: usize = base_batches.iter().map(|b| b.num_rows()).sum();
let delta_row_count: usize = delta_batches.iter().map(|b| b.num_rows()).sum();
let has_tombstones = !tombstones.is_empty();
let nothing_to_do = delta_batches.is_empty() && !has_tombstones;
if nothing_to_do {
return Ok(CompactResult {
base_rows: base_row_count,
delta_rows: 0,
final_rows: base_row_count,
deltas_merged: 0,
tombstones_applied: 0,
rows_dropped_by_tombstones: 0,
});
}
base_batches.extend(delta_batches);
let pre_filter_rows: usize = base_row_count + delta_row_count;
// If primary key specified, deduplicate (keep last occurrence)
let merged_batches = if let Some(_pk) = primary_key_col {
// Current simplification — full dedup requires sort by PK; leave
// that work for a follow-up. Tombstone filter still runs below.
base_batches
} else {
base_batches
};
// Tombstone filter: drop rows whose row_key matches.
let (final_batches, dropped) = if has_tombstones {
apply_tombstone_filter(merged_batches, tombstones)?
} else {
(merged_batches, 0)
};
let final_rows: usize = final_batches.iter().map(|b| b.num_rows()).sum();
// Write merged base atomically: build ONE Parquet file (single
// footer!) via ArrowWriter, stage it under a temporary key, read
// back to verify row count matches what we intended, THEN swap
// into base_key. If anything goes wrong between build and swap,
// the original base_key is untouched.
//
// Why this matters: concatenating N record_batch_to_parquet()
// outputs produces N independent Parquet files glued together —
// readers see only the first footer. That silently corrupted the
// candidates table before this code was fixed.
let schema_ref = final_batches
.first()
.map(|b| b.schema())
.ok_or_else(|| "no batches to write".to_string())?;
// Use Snappy to match the compression of files written by ingest —
// otherwise compaction would inflate file size 3× on each rewrite.
let props = WriterProperties::builder()
.set_compression(Compression::SNAPPY)
.build();
let mut buf: Vec<u8> = Vec::with_capacity(16 * 1024 * 1024);
{
let mut writer = ArrowWriter::try_new(&mut buf, schema_ref.clone(), Some(props))
.map_err(|e| format!("ArrowWriter init: {e}"))?;
for batch in &final_batches {
writer.write(batch).map_err(|e| format!("write batch: {e}"))?;
}
writer.close().map_err(|e| format!("close writer: {e}"))?;
}
// Verify before committing — if the written bytes don't parse as
// the row count we expect, DON'T overwrite the base file.
let (_, verify_batches) = parquet_to_record_batches(&buf)
.map_err(|e| format!("verify parse: {e}"))?;
let verify_rows: usize = verify_batches.iter().map(|b| b.num_rows()).sum();
if verify_rows != final_rows {
return Err(format!(
"compact verification failed: wrote {} bytes that parse as {} rows (expected {}); base_key '{}' untouched",
buf.len(), verify_rows, final_rows, base_key,
));
}
// Stage under a temp key first, then promote via a second put
// over base_key. Object stores don't have atomic rename across
// arbitrary backends; this is "write twice + delete temp" which
// is the widely-adopted pattern.
let temp_key = format!("{base_key}.compact-{}.tmp", chrono::Utc::now().timestamp_millis());
ops::put(store, &temp_key, Bytes::from(buf.clone())).await?;
ops::put(store, base_key, Bytes::from(buf)).await?;
let _ = ops::delete(store, &temp_key).await;
// Only now that base is durably updated, delete delta files.
let delta_keys = list_deltas(store, dataset_name).await?;
for key in &delta_keys {
let _ = ops::delete(store, key).await;
}
tracing::info!(
"compacted '{}': {} deltas merged, {} tombstones applied ({} rows dropped), {} → {} rows",
dataset_name, delta_count, tombstones.len(), dropped, pre_filter_rows, final_rows,
);
Ok(CompactResult {
base_rows: base_row_count,
delta_rows: delta_row_count,
final_rows,
deltas_merged: delta_count,
tombstones_applied: tombstones.len(),
rows_dropped_by_tombstones: dropped,
})
}
/// Filter every batch in `batches` to exclude rows whose `row_key_column`
/// value appears in the tombstone set. All tombstones must share one
/// row_key_column — enforced at tombstone write time.
fn apply_tombstone_filter(
batches: Vec<RecordBatch>,
tombstones: &[Tombstone],
) -> Result<(Vec<RecordBatch>, usize), String> {
let col_name = &tombstones[0].row_key_column;
let bad_values: HashSet<String> = tombstones
.iter()
.map(|t| t.row_key_value.clone())
.collect();
let mut out = Vec::with_capacity(batches.len());
let mut total_dropped = 0usize;
for batch in batches {
let col_idx = match batch.schema().index_of(col_name) {
Ok(i) => i,
Err(_) => {
// Column not in this batch — nothing to drop, pass through.
out.push(batch);
continue;
}
};
let col = batch.column(col_idx);
// Build a "keep" mask: true where row is NOT tombstoned.
let mask: Vec<bool> = if let Some(s) = col.as_any().downcast_ref::<StringArray>() {
(0..s.len()).map(|i| {
s.is_null(i) || !bad_values.contains(s.value(i))
}).collect()
} else if let Some(a) = col.as_any().downcast_ref::<Int32Array>() {
(0..a.len()).map(|i| {
a.is_null(i) || !bad_values.contains(&a.value(i).to_string())
}).collect()
} else if let Some(a) = col.as_any().downcast_ref::<Int64Array>() {
(0..a.len()).map(|i| {
a.is_null(i) || !bad_values.contains(&a.value(i).to_string())
}).collect()
} else {
// Key column of an unsupported type — log and pass through
// rather than panic. Operator has to handle this manually.
tracing::warn!(
"tombstone filter: column '{}' has unsupported type {:?}; passing rows through",
col_name, col.data_type(),
);
out.push(batch);
continue;
};
let kept = mask.iter().filter(|b| **b).count();
let dropped = mask.len() - kept;
total_dropped += dropped;
if dropped == 0 {
out.push(batch);
continue;
}
let bool_arr = BooleanArray::from(mask);
let filtered = filter_record_batch(&batch, &bool_arr)
.map_err(|e| format!("filter batch: {e}"))?;
out.push(filtered);
}
Ok((out, total_dropped))
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct CompactResult {
pub base_rows: usize,
pub delta_rows: usize,
pub final_rows: usize,
pub deltas_merged: usize,
pub tombstones_applied: usize,
pub rows_dropped_by_tombstones: usize,
}
+7
View File
@@ -0,0 +1,7 @@
pub mod cache;
pub mod context;
pub mod delta;
pub mod paged;
pub mod service;
pub mod workspace;
pub mod workspace_service;
+222
View File
@@ -0,0 +1,222 @@
/// Paged query results with server-side cursor.
/// Large results are stored server-side, client fetches pages on demand.
/// Pattern matches the embedding supervisor: chunk → process → retry.
use arrow::array::RecordBatch;
use arrow::json::writer::{JsonArray, Writer as JsonWriter};
use serde::Serialize;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use crate::context::QueryEngine;
/// A stored query result that clients page through.
#[derive(Clone)]
pub struct StoredResult {
pub query_id: String,
pub sql: String,
pub columns: Vec<ColumnInfo>,
pub total_rows: usize,
pub batches: Vec<RecordBatch>,
pub created_at: std::time::Instant,
}
#[derive(Clone, Serialize)]
pub struct ColumnInfo {
pub name: String,
pub data_type: String,
}
/// Response for initial query — metadata only, no rows.
#[derive(Serialize)]
pub struct QueryHandle {
pub query_id: String,
pub total_rows: usize,
pub columns: Vec<ColumnInfo>,
pub total_pages: usize,
pub page_size: usize,
}
/// A page of results.
#[derive(Serialize)]
pub struct PageResult {
pub query_id: String,
pub page: usize,
pub page_size: usize,
pub total_rows: usize,
pub total_pages: usize,
pub rows: serde_json::Value,
pub row_count: usize,
}
/// Server-side result store with TTL.
#[derive(Clone)]
pub struct ResultStore {
results: Arc<RwLock<HashMap<String, StoredResult>>>,
default_page_size: usize,
max_results: usize,
}
impl ResultStore {
pub fn new(default_page_size: usize, max_results: usize) -> Self {
Self {
results: Arc::new(RwLock::new(HashMap::new())),
default_page_size,
max_results,
}
}
/// Execute query and store results. Returns handle with metadata.
pub async fn execute_and_store(
&self,
engine: &QueryEngine,
sql: &str,
) -> Result<QueryHandle, String> {
let batches = engine.query(sql).await?;
if batches.is_empty() {
let qid = format!("qr-{}", chrono::Utc::now().timestamp_millis());
return Ok(QueryHandle {
query_id: qid,
total_rows: 0,
columns: vec![],
total_pages: 0,
page_size: self.default_page_size,
});
}
let schema = batches[0].schema();
let columns: Vec<ColumnInfo> = schema.fields().iter().map(|f| ColumnInfo {
name: f.name().clone(),
data_type: f.data_type().to_string(),
}).collect();
let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
let total_pages = (total_rows + self.default_page_size - 1) / self.default_page_size;
let qid = format!("qr-{}", chrono::Utc::now().timestamp_millis());
let result = StoredResult {
query_id: qid.clone(),
sql: sql.to_string(),
columns: columns.clone(),
total_rows,
batches,
created_at: std::time::Instant::now(),
};
// Store (evict old results if at capacity)
let mut store = self.results.write().await;
if store.len() >= self.max_results {
// Evict oldest
if let Some(oldest_key) = store.iter()
.min_by_key(|(_, v)| v.created_at)
.map(|(k, _)| k.clone())
{
store.remove(&oldest_key);
}
}
store.insert(qid.clone(), result);
Ok(QueryHandle {
query_id: qid,
total_rows,
columns,
total_pages,
page_size: self.default_page_size,
})
}
/// Get a page of results. Handles chunking and retry-safe.
pub async fn get_page(
&self,
query_id: &str,
page: usize,
page_size: Option<usize>,
) -> Result<PageResult, String> {
let store = self.results.read().await;
let result = store.get(query_id)
.ok_or_else(|| format!("query result not found: {query_id} (may have expired)"))?;
let ps = page_size.unwrap_or(self.default_page_size);
let start_row = page * ps;
if start_row >= result.total_rows {
return Ok(PageResult {
query_id: query_id.to_string(),
page,
page_size: ps,
total_rows: result.total_rows,
total_pages: (result.total_rows + ps - 1) / ps,
rows: serde_json::Value::Array(vec![]),
row_count: 0,
});
}
let end_row = (start_row + ps).min(result.total_rows);
// Extract the right rows from batches
let page_batches = slice_batches(&result.batches, start_row, end_row);
let rows = batches_to_json(&page_batches)?;
let row_count = rows.as_array().map(|a| a.len()).unwrap_or(0);
Ok(PageResult {
query_id: query_id.to_string(),
page,
page_size: ps,
total_rows: result.total_rows,
total_pages: (result.total_rows + ps - 1) / ps,
rows,
row_count,
})
}
pub fn default_page_size(&self) -> usize {
self.default_page_size
}
}
/// Slice record batches to extract rows [start, end).
fn slice_batches(batches: &[RecordBatch], start: usize, end: usize) -> Vec<RecordBatch> {
let mut result = Vec::new();
let mut current_offset = 0;
for batch in batches {
let batch_end = current_offset + batch.num_rows();
if batch_end <= start {
current_offset = batch_end;
continue;
}
if current_offset >= end {
break;
}
let local_start = if start > current_offset { start - current_offset } else { 0 };
let local_end = if end < batch_end { end - current_offset } else { batch.num_rows() };
let length = local_end - local_start;
if length > 0 {
result.push(batch.slice(local_start, length));
}
current_offset = batch_end;
}
result
}
fn batches_to_json(batches: &[RecordBatch]) -> Result<serde_json::Value, String> {
if batches.is_empty() {
return Ok(serde_json::Value::Array(vec![]));
}
let mut buf = Vec::new();
let mut writer = JsonWriter::<_, JsonArray>::new(&mut buf);
for batch in batches {
writer.write(batch).map_err(|e| format!("JSON write: {e}"))?;
}
writer.finish().map_err(|e| format!("JSON finish: {e}"))?;
drop(writer);
serde_json::from_slice(&buf).map_err(|e| format!("JSON parse: {e}"))
}
+332
View File
@@ -0,0 +1,332 @@
use arrow::array::RecordBatch;
use arrow::json::writer::{JsonArray, Writer as JsonWriter};
use axum::{
Json, Router,
extract::{Path, Query, State},
http::StatusCode,
response::IntoResponse,
routing::{get, post},
};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use truth::{RuleAction, TruthStore};
use crate::context::QueryEngine;
use crate::delta;
use crate::paged::ResultStore;
#[derive(Clone)]
pub struct QueryState {
pub engine: QueryEngine,
pub result_store: ResultStore,
// Policy gate for incoming SQL. Every /sql and /paged request is
// evaluated against this store before hitting DataFusion. Added for
// P42-002 ("raw SQL forwarded without schema or policy gate") after
// the scrum master's queryd/service.rs finding looped across iters
// 3-5 without ever being reachable by the 6-line auto-applier.
pub truth: Arc<TruthStore>,
}
pub fn router(engine: QueryEngine) -> Router {
router_with_truth(engine, Arc::new(truth::sql_query_guard_store()))
}
/// Test/integration hook: construct the router with a caller-supplied
/// TruthStore so tests can assert reject/pass behavior deterministically
/// without depending on the default needle list.
pub fn router_with_truth(engine: QueryEngine, truth: Arc<TruthStore>) -> Router {
let state = QueryState {
engine: engine.clone(),
result_store: ResultStore::new(100, 50), // 100 rows/page, keep 50 results
truth,
};
Router::new()
.route("/health", get(health))
.route("/sql", post(execute_query))
.route("/paged", post(paged_query))
.route("/page/{query_id}/{page}", get(get_page))
.route("/cache/pin", post(pin_dataset))
.route("/cache/evict", post(evict_dataset))
.route("/cache/stats", get(cache_stats))
.route("/compact", post(compact_dataset))
.with_state(state)
}
async fn health() -> &'static str {
"queryd ok"
}
// --- SQL Query ---
#[derive(Deserialize)]
struct QueryRequest {
sql: String,
}
#[derive(Serialize)]
struct QueryResponse {
columns: Vec<ColumnInfo>,
rows: serde_json::Value,
row_count: usize,
// Elapsed wall time from handler entry to response. Required for
// audit-log parity — gateway's audit row previously stored null here.
// Scrum iter 9 finding, populated from std::time::Instant captured
// at the top of execute_query / paged_query.
latency_ms: u64,
}
#[derive(Serialize)]
struct ColumnInfo {
name: String,
data_type: String,
}
fn batches_to_json(batches: &[RecordBatch]) -> Result<serde_json::Value, String> {
let mut buf = Vec::new();
let mut writer = JsonWriter::<_, JsonArray>::new(&mut buf);
for batch in batches {
writer.write(batch).map_err(|e| format!("JSON write error: {e}"))?;
}
writer.finish().map_err(|e| format!("JSON finish error: {e}"))?;
drop(writer);
serde_json::from_slice(&buf).map_err(|e| format!("JSON parse error: {e}"))
}
/// Evaluate the request SQL against the configured TruthStore. Returns
/// the Reject/Block message on the first failing mandatory rule so the
/// handler can short-circuit. Returns None when all rules pass (or when
/// the failures' declared action is non-mandatory like Redact/Pass).
fn sql_policy_check(truth: &TruthStore, sql: &str) -> Option<String> {
let ctx = serde_json::json!({ "sql": sql });
for outcome in truth.evaluate("sql_query", &ctx) {
if !outcome.passed {
// FieldEmpty / FieldContainsAny etc. are enforced only when
// condition HOLDS (i.e. passed=true). Below means "passed=false",
// so the rule condition did not hold — no enforcement.
continue;
}
match &outcome.action {
RuleAction::Reject { message } | RuleAction::Block { message } => {
return Some(message.clone());
}
_ => {}
}
}
None
}
async fn execute_query(
State(state): State<QueryState>,
Json(req): Json<QueryRequest>,
) -> impl IntoResponse {
let started = std::time::Instant::now();
tracing::info!("executing query: {}", req.sql);
if let Some(reason) = sql_policy_check(&state.truth, &req.sql) {
tracing::warn!("sql rejected by truth gate: {reason}");
return Err((StatusCode::FORBIDDEN, reason));
}
match state.engine.query(&req.sql).await {
Ok(batches) => {
if batches.is_empty() {
return Ok(Json(QueryResponse {
columns: vec![],
rows: serde_json::Value::Array(vec![]),
row_count: 0,
latency_ms: started.elapsed().as_millis() as u64,
}));
}
let schema = batches[0].schema();
let columns: Vec<ColumnInfo> = schema.fields().iter().map(|f| ColumnInfo {
name: f.name().clone(),
data_type: f.data_type().to_string(),
}).collect();
let rows = batches_to_json(&batches)
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
let row_count = rows.as_array().map(|a| a.len()).unwrap_or(0);
Ok(Json(QueryResponse {
columns,
rows,
row_count,
latency_ms: started.elapsed().as_millis() as u64,
}))
}
Err(e) => Err((StatusCode::BAD_REQUEST, e)),
}
}
// --- Paged Queries (large result sets) ---
async fn paged_query(
State(state): State<QueryState>,
Json(req): Json<QueryRequest>,
) -> impl IntoResponse {
tracing::info!("paged query: {}", req.sql);
if let Some(reason) = sql_policy_check(&state.truth, &req.sql) {
tracing::warn!("paged sql rejected by truth gate: {reason}");
return Err((StatusCode::FORBIDDEN, reason));
}
match state.result_store.execute_and_store(&state.engine, &req.sql).await {
Ok(handle) => Ok(Json(handle)),
Err(e) => Err((StatusCode::BAD_REQUEST, e)),
}
}
#[derive(Deserialize)]
struct PageQuery {
size: Option<usize>,
}
async fn get_page(
State(state): State<QueryState>,
Path((query_id, page)): Path<(String, usize)>,
Query(q): Query<PageQuery>,
) -> impl IntoResponse {
match state.result_store.get_page(&query_id, page, q.size).await {
Ok(result) => Ok(Json(result)),
Err(e) => Err((StatusCode::NOT_FOUND, e)),
}
}
// --- Cache Management ---
#[derive(Deserialize)]
struct CacheRequest {
dataset: String,
}
async fn pin_dataset(
State(state): State<QueryState>,
Json(req): Json<CacheRequest>,
) -> impl IntoResponse {
match state.engine.pin_dataset(&req.dataset).await {
Ok(()) => Ok((StatusCode::OK, format!("pinned: {}", req.dataset))),
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
async fn evict_dataset(
State(state): State<QueryState>,
Json(req): Json<CacheRequest>,
) -> impl IntoResponse {
if state.engine.cache().evict(&req.dataset).await {
(StatusCode::OK, format!("evicted: {}", req.dataset))
} else {
(StatusCode::NOT_FOUND, format!("not cached: {}", req.dataset))
}
}
async fn cache_stats(State(state): State<QueryState>) -> impl IntoResponse {
let stats = state.engine.cache().stats().await;
Json(stats)
}
// --- Compaction ---
#[derive(Deserialize)]
struct CompactRequest {
dataset: String,
base_key: String,
primary_key: Option<String>,
}
async fn compact_dataset(
State(state): State<QueryState>,
Json(req): Json<CompactRequest>,
) -> impl IntoResponse {
// Phase E: pull tombstones for this dataset and let compact physically
// drop those rows. After a successful rewrite, clear the tombstone log
// — the rows are gone from disk, the tombstones have done their job.
let tombstones = state
.engine
.catalog()
.list_tombstones(&req.dataset)
.await
.unwrap_or_default();
match delta::compact(
state.engine.store(),
&req.dataset,
&req.base_key,
req.primary_key.as_deref(),
&tombstones,
).await {
Ok(result) => {
if result.rows_dropped_by_tombstones > 0 {
if let Err(e) = state.engine.catalog().tombstones().clear(&req.dataset).await {
tracing::warn!("post-compact tombstone clear failed: {e}");
}
}
Ok(Json(result))
}
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
#[cfg(test)]
mod sql_policy_tests {
use super::*;
use truth::sql_query_guard_store;
// These tests exercise the policy gate without spinning up a DataFusion
// engine — they only need `TruthStore`. Purpose: prove the P42-002
// enforcement point actually rejects destructive SQL. This is the
// regression guard for the queryd/service.rs finding that looped
// across scrum iters 3-5.
#[test]
fn blocks_drop_table() {
let store = sql_query_guard_store();
let reason = sql_policy_check(&store, "DROP TABLE users").expect("must reject");
assert!(reason.contains("destructive"), "reason: {reason}");
}
#[test]
fn blocks_delete_from() {
let store = sql_query_guard_store();
assert!(sql_policy_check(&store, "delete from t where 1=1").is_some());
}
#[test]
fn blocks_truncate() {
let store = sql_query_guard_store();
assert!(sql_policy_check(&store, "TRUNCATE workers").is_some());
}
#[test]
fn blocks_empty_sql() {
let store = sql_query_guard_store();
assert!(sql_policy_check(&store, "").is_some());
}
#[test]
fn allows_benign_select() {
let store = sql_query_guard_store();
assert!(sql_policy_check(&store, "SELECT count(*) FROM workers").is_none());
}
#[test]
fn allows_select_with_deleted_word_in_column() {
// Substring match is narrow ("delete from", not "delete"), so a
// column named `deleted_at` doesn't trip the guard. Important
// check — false positives on benign queries would make the gate
// unusable in practice.
let store = sql_query_guard_store();
assert!(
sql_policy_check(&store, "SELECT deleted_at FROM t").is_none(),
"column names containing 'delete' must not be rejected"
);
}
#[test]
fn case_insensitive_match_catches_mixed_case() {
let store = sql_query_guard_store();
assert!(sql_policy_check(&store, "Drop Table X").is_some());
}
}
+269
View File
@@ -0,0 +1,269 @@
/// Agent workspaces — named overlays for contract/search-specific work.
/// Each workspace tracks an agent's activity on a specific contract or search,
/// with daily/weekly/monthly tiers and instant handoff capability.
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use object_store::ObjectStore;
/// Retention tier for workspace data.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Tier {
Daily, // expires end of day, active search scratch
Weekly, // expires end of week, active contract work
Monthly, // expires end of month, contract lifecycle
Pinned, // never expires, manually managed
}
/// A saved query/filter within a workspace.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SavedSearch {
pub name: String,
pub sql: String,
pub created_at: DateTime<Utc>,
}
/// A shortlisted candidate or record.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ShortlistEntry {
pub dataset: String,
pub record_id: String,
pub notes: String,
pub added_at: DateTime<Utc>,
pub added_by: String,
}
/// Activity log entry.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ActivityEntry {
pub action: String, // "search", "shortlist", "call", "email", "update", "ingest"
pub detail: String,
pub timestamp: DateTime<Utc>,
pub agent: String,
}
/// A workspace — an agent's working context for a contract or search.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Workspace {
pub id: String,
pub name: String, // "Apex Corp .NET Developers - Chicago"
pub description: String,
pub tier: Tier,
pub owner: String, // current agent
pub previous_owners: Vec<HandoffRecord>,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
// Work content
pub saved_searches: Vec<SavedSearch>,
pub shortlist: Vec<ShortlistEntry>,
pub activity: Vec<ActivityEntry>,
pub ingested_datasets: Vec<String>, // datasets this workspace created
pub delta_keys: Vec<String>, // delta files specific to this workspace
pub tags: Vec<String>,
}
/// Record of a handoff between agents.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HandoffRecord {
pub from_agent: String,
pub to_agent: String,
pub reason: String,
pub timestamp: DateTime<Utc>,
}
/// Workspace manager — in-memory registry with persistence.
#[derive(Clone)]
pub struct WorkspaceManager {
workspaces: Arc<RwLock<HashMap<String, Workspace>>>,
store: Arc<dyn ObjectStore>,
}
impl WorkspaceManager {
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
Self {
workspaces: Arc::new(RwLock::new(HashMap::new())),
store,
}
}
/// Rebuild from persisted workspace files on startup.
pub async fn rebuild(&self) -> Result<usize, String> {
let keys = storaged::ops::list(&self.store, Some("workspaces/")).await?;
let mut ws_map = self.workspaces.write().await;
ws_map.clear();
for key in &keys {
if !key.ends_with(".json") { continue; }
let data = storaged::ops::get(&self.store, key).await?;
match serde_json::from_slice::<Workspace>(&data) {
Ok(ws) => { ws_map.insert(ws.id.clone(), ws); }
Err(e) => tracing::warn!("failed to load workspace {key}: {e}"),
}
}
let count = ws_map.len();
if count > 0 {
tracing::info!("loaded {count} workspaces");
}
Ok(count)
}
/// Create a new workspace.
pub async fn create(&self, name: String, description: String, owner: String, tier: Tier) -> Result<Workspace, String> {
let now = Utc::now();
let id = format!("ws-{}", now.timestamp_millis());
let ws = Workspace {
id: id.clone(),
name,
description,
tier,
owner,
previous_owners: vec![],
created_at: now,
updated_at: now,
saved_searches: vec![],
shortlist: vec![],
activity: vec![],
ingested_datasets: vec![],
delta_keys: vec![],
tags: vec![],
};
self.persist(&ws).await?;
self.workspaces.write().await.insert(id.clone(), ws.clone());
tracing::info!("created workspace: {} ({})", ws.name, ws.id);
Ok(ws)
}
/// Handoff workspace to another agent. Instant — no data copy.
pub async fn handoff(&self, workspace_id: &str, to_agent: &str, reason: &str) -> Result<Workspace, String> {
let mut ws_map = self.workspaces.write().await;
let ws = ws_map.get_mut(workspace_id)
.ok_or_else(|| format!("workspace not found: {workspace_id}"))?;
let record = HandoffRecord {
from_agent: ws.owner.clone(),
to_agent: to_agent.to_string(),
reason: reason.to_string(),
timestamp: Utc::now(),
};
ws.previous_owners.push(record);
ws.owner = to_agent.to_string();
ws.updated_at = Utc::now();
ws.activity.push(ActivityEntry {
action: "handoff".to_string(),
detail: format!("handed off to {} — {}", to_agent, reason),
timestamp: Utc::now(),
agent: to_agent.to_string(),
});
let ws_clone = ws.clone();
drop(ws_map);
self.persist(&ws_clone).await?;
tracing::info!("workspace '{}' handed off to {}", ws_clone.name, to_agent);
Ok(ws_clone)
}
/// Add a saved search to a workspace.
pub async fn add_search(&self, workspace_id: &str, name: String, sql: String, agent: &str) -> Result<(), String> {
let mut ws_map = self.workspaces.write().await;
let ws = ws_map.get_mut(workspace_id)
.ok_or_else(|| format!("workspace not found: {workspace_id}"))?;
ws.saved_searches.push(SavedSearch {
name: name.clone(),
sql,
created_at: Utc::now(),
});
ws.activity.push(ActivityEntry {
action: "search".into(),
detail: format!("saved search: {name}"),
timestamp: Utc::now(),
agent: agent.to_string(),
});
ws.updated_at = Utc::now();
let ws_clone = ws.clone();
drop(ws_map);
self.persist(&ws_clone).await
}
/// Add a candidate/record to the shortlist.
pub async fn add_to_shortlist(&self, workspace_id: &str, dataset: String, record_id: String, notes: String, agent: &str) -> Result<(), String> {
let mut ws_map = self.workspaces.write().await;
let ws = ws_map.get_mut(workspace_id)
.ok_or_else(|| format!("workspace not found: {workspace_id}"))?;
ws.shortlist.push(ShortlistEntry {
dataset: dataset.clone(),
record_id: record_id.clone(),
notes,
added_at: Utc::now(),
added_by: agent.to_string(),
});
ws.activity.push(ActivityEntry {
action: "shortlist".into(),
detail: format!("added {record_id} from {dataset}"),
timestamp: Utc::now(),
agent: agent.to_string(),
});
ws.updated_at = Utc::now();
let ws_clone = ws.clone();
drop(ws_map);
self.persist(&ws_clone).await
}
/// Log an activity.
pub async fn log_activity(&self, workspace_id: &str, action: String, detail: String, agent: &str) -> Result<(), String> {
let mut ws_map = self.workspaces.write().await;
let ws = ws_map.get_mut(workspace_id)
.ok_or_else(|| format!("workspace not found: {workspace_id}"))?;
ws.activity.push(ActivityEntry {
action,
detail,
timestamp: Utc::now(),
agent: agent.to_string(),
});
ws.updated_at = Utc::now();
let ws_clone = ws.clone();
drop(ws_map);
self.persist(&ws_clone).await
}
/// Get a workspace.
pub async fn get(&self, workspace_id: &str) -> Option<Workspace> {
self.workspaces.read().await.get(workspace_id).cloned()
}
/// List all workspaces, optionally filtered by owner or tier.
pub async fn list(&self, owner: Option<&str>, tier: Option<&Tier>) -> Vec<Workspace> {
let ws_map = self.workspaces.read().await;
ws_map.values()
.filter(|ws| {
owner.map_or(true, |o| ws.owner == o) &&
tier.map_or(true, |t| ws.tier == *t)
})
.cloned()
.collect()
}
/// Persist workspace to object storage.
async fn persist(&self, ws: &Workspace) -> Result<(), String> {
let key = format!("workspaces/{}.json", ws.id);
let json = serde_json::to_vec_pretty(ws).map_err(|e| e.to_string())?;
storaged::ops::put(&self.store, &key, json.into()).await
}
}
+136
View File
@@ -0,0 +1,136 @@
use axum::{
Json, Router,
extract::{Path, Query, State},
http::StatusCode,
response::IntoResponse,
routing::{get, post},
};
use serde::Deserialize;
use crate::workspace::{Tier, WorkspaceManager};
pub fn router(manager: WorkspaceManager) -> Router {
Router::new()
.route("/", get(list_workspaces))
.route("/create", post(create_workspace))
.route("/{id}", get(get_workspace))
.route("/{id}/handoff", post(handoff))
.route("/{id}/search", post(add_search))
.route("/{id}/shortlist", post(add_to_shortlist))
.route("/{id}/activity", post(log_activity))
.with_state(manager)
}
#[derive(Deserialize)]
struct ListQuery {
owner: Option<String>,
tier: Option<Tier>,
}
async fn list_workspaces(
State(mgr): State<WorkspaceManager>,
Query(q): Query<ListQuery>,
) -> impl IntoResponse {
let workspaces = mgr.list(q.owner.as_deref(), q.tier.as_ref()).await;
Json(workspaces)
}
#[derive(Deserialize)]
struct CreateRequest {
name: String,
description: String,
owner: String,
tier: Tier,
}
async fn create_workspace(
State(mgr): State<WorkspaceManager>,
Json(req): Json<CreateRequest>,
) -> impl IntoResponse {
match mgr.create(req.name, req.description, req.owner, req.tier).await {
Ok(ws) => Ok((StatusCode::CREATED, Json(ws))),
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
async fn get_workspace(
State(mgr): State<WorkspaceManager>,
Path(id): Path<String>,
) -> impl IntoResponse {
match mgr.get(&id).await {
Some(ws) => Ok(Json(ws)),
None => Err((StatusCode::NOT_FOUND, format!("workspace not found: {id}"))),
}
}
#[derive(Deserialize)]
struct HandoffRequest {
to_agent: String,
reason: String,
}
async fn handoff(
State(mgr): State<WorkspaceManager>,
Path(id): Path<String>,
Json(req): Json<HandoffRequest>,
) -> impl IntoResponse {
match mgr.handoff(&id, &req.to_agent, &req.reason).await {
Ok(ws) => Ok(Json(ws)),
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
#[derive(Deserialize)]
struct SearchRequest {
name: String,
sql: String,
agent: String,
}
async fn add_search(
State(mgr): State<WorkspaceManager>,
Path(id): Path<String>,
Json(req): Json<SearchRequest>,
) -> impl IntoResponse {
match mgr.add_search(&id, req.name, req.sql, &req.agent).await {
Ok(()) => Ok((StatusCode::OK, "search saved")),
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
#[derive(Deserialize)]
struct ShortlistRequest {
dataset: String,
record_id: String,
notes: String,
agent: String,
}
async fn add_to_shortlist(
State(mgr): State<WorkspaceManager>,
Path(id): Path<String>,
Json(req): Json<ShortlistRequest>,
) -> impl IntoResponse {
match mgr.add_to_shortlist(&id, req.dataset, req.record_id, req.notes, &req.agent).await {
Ok(()) => Ok((StatusCode::OK, "added to shortlist")),
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
#[derive(Deserialize)]
struct ActivityRequest {
action: String,
detail: String,
agent: String,
}
async fn log_activity(
State(mgr): State<WorkspaceManager>,
Path(id): Path<String>,
Json(req): Json<ActivityRequest>,
) -> impl IntoResponse {
match mgr.log_activity(&id, req.action, req.detail, &req.agent).await {
Ok(()) => Ok((StatusCode::OK, "activity logged")),
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
+19
View File
@@ -0,0 +1,19 @@
[package]
name = "shared"
version = "0.1.0"
edition = "2024"
[dependencies]
serde = { workspace = true }
serde_json = { workspace = true }
thiserror = { workspace = true }
uuid = { workspace = true }
chrono = { workspace = true }
arrow = { workspace = true }
parquet = { workspace = true }
bytes = { workspace = true }
sha2 = { workspace = true }
toml = { workspace = true }
tracing = { workspace = true }
tokio = { workspace = true }
async-trait = "0.1"
+83
View File
@@ -0,0 +1,83 @@
use arrow::array::RecordBatch;
use arrow::datatypes::SchemaRef;
use bytes::Bytes;
use parquet::arrow::ArrowWriter;
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use sha2::{Digest, Sha256};
use crate::types::SchemaFingerprint;
/// Write a RecordBatch to Parquet bytes.
pub fn record_batch_to_parquet(batch: &RecordBatch) -> Result<Bytes, String> {
let mut buf = Vec::new();
let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), None)
.map_err(|e| e.to_string())?;
writer.write(batch).map_err(|e| e.to_string())?;
writer.close().map_err(|e| e.to_string())?;
Ok(Bytes::from(buf))
}
/// Read Parquet bytes into a Vec of RecordBatches.
pub fn parquet_to_record_batches(data: &[u8]) -> Result<(SchemaRef, Vec<RecordBatch>), String> {
let builder = ParquetRecordBatchReaderBuilder::try_new(Bytes::copy_from_slice(data))
.map_err(|e| e.to_string())?;
let schema = builder.schema().clone();
let reader = builder.build().map_err(|e| e.to_string())?;
let batches: Vec<RecordBatch> = reader.collect::<Result<Vec<_>, _>>().map_err(|e| e.to_string())?;
Ok((schema, batches))
}
/// Compute a deterministic fingerprint from an Arrow schema.
pub fn fingerprint_schema(schema: &SchemaRef) -> SchemaFingerprint {
let mut hasher = Sha256::new();
for field in schema.fields() {
hasher.update(field.name().as_bytes());
hasher.update(field.data_type().to_string().as_bytes());
hasher.update(if field.is_nullable() { b"1" } else { b"0" });
}
let hash = hasher.finalize();
SchemaFingerprint(format!("{hash:x}"))
}
#[cfg(test)]
mod tests {
use super::*;
use arrow::array::{Float64Array, Int32Array, StringArray};
use arrow::datatypes::{DataType, Field, Schema};
use std::sync::Arc;
fn sample_batch() -> RecordBatch {
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int32, false),
Field::new("name", DataType::Utf8, false),
Field::new("score", DataType::Float64, true),
]));
RecordBatch::try_new(
schema,
vec![
Arc::new(Int32Array::from(vec![1, 2, 3])),
Arc::new(StringArray::from(vec!["alice", "bob", "carol"])),
Arc::new(Float64Array::from(vec![9.5, 8.2, 7.8])),
],
)
.unwrap()
}
#[test]
fn roundtrip_parquet() {
let batch = sample_batch();
let parquet_bytes = record_batch_to_parquet(&batch).unwrap();
let (schema, batches) = parquet_to_record_batches(&parquet_bytes).unwrap();
assert_eq!(schema.fields().len(), 3);
assert_eq!(batches.len(), 1);
assert_eq!(batches[0].num_rows(), 3);
}
#[test]
fn schema_fingerprint_deterministic() {
let batch = sample_batch();
let fp1 = fingerprint_schema(&batch.schema());
let fp2 = fingerprint_schema(&batch.schema());
assert_eq!(fp1, fp2);
}
}
+203
View File
@@ -0,0 +1,203 @@
use serde::Deserialize;
use std::path::Path;
#[derive(Debug, Clone, Deserialize)]
pub struct Config {
pub gateway: GatewayConfig,
pub storage: StorageConfig,
#[serde(default)]
pub catalog: CatalogConfig,
#[serde(default)]
pub query: QueryConfig,
pub sidecar: SidecarConfig,
#[serde(default)]
pub ai: AiConfig,
#[serde(default)]
pub auth: AuthConfig,
#[serde(default)]
pub observability: ObservabilityConfig,
#[serde(default)]
pub agent: AgentSettings,
}
/// Phase 16.2 — background autotune agent settings.
///
/// Duplicated from `vectord::agent::AgentConfig` because `shared` can't
/// depend on `vectord` (vectord already depends on shared). The gateway
/// copies these into the vectord config at startup.
#[derive(Debug, Clone, Deserialize)]
pub struct AgentSettings {
#[serde(default)]
pub enabled: bool,
#[serde(default = "default_cycle_interval_secs")]
pub cycle_interval_secs: u64,
#[serde(default = "default_cooldown_secs")]
pub cooldown_between_trials_secs: u64,
#[serde(default = "default_min_recall")]
pub min_recall: f32,
#[serde(default = "default_max_trials_per_hour")]
pub max_trials_per_hour: u32,
}
impl Default for AgentSettings {
fn default() -> Self {
Self {
enabled: false,
cycle_interval_secs: default_cycle_interval_secs(),
cooldown_between_trials_secs: default_cooldown_secs(),
min_recall: default_min_recall(),
max_trials_per_hour: default_max_trials_per_hour(),
}
}
}
fn default_cycle_interval_secs() -> u64 { 60 }
fn default_cooldown_secs() -> u64 { 30 }
fn default_min_recall() -> f32 { 0.9 }
fn default_max_trials_per_hour() -> u32 { 30 }
#[derive(Debug, Clone, Deserialize)]
pub struct GatewayConfig {
#[serde(default = "default_host")]
pub host: String,
#[serde(default = "default_gateway_port")]
pub port: u16,
}
#[derive(Debug, Clone, Deserialize)]
pub struct StorageConfig {
/// Legacy single-backend root. If `buckets` is empty, this is used to
/// create an implicit `primary` bucket at this path — preserves the
/// pre-federation config shape.
#[serde(default = "default_storage_root")]
pub root: String,
/// Where profile buckets are rooted when auto-provisioned.
#[serde(default = "default_profile_root")]
pub profile_root: String,
/// Name of the bucket used for read fallback when a target bucket is
/// unreachable. If `None`, no fallback — reads fail hard.
#[serde(default)]
pub rescue_bucket: Option<String>,
/// Explicitly configured buckets. Empty = backward-compat single-bucket
/// mode driven by `root`.
#[serde(default)]
pub buckets: Vec<BucketConfig>,
}
#[derive(Debug, Clone, Deserialize)]
pub struct BucketConfig {
pub name: String,
pub backend: String, // "local" | "s3"
/// Local filesystem root (for backend = "local")
pub root: Option<String>,
/// S3 bucket name (for backend = "s3")
pub bucket: Option<String>,
pub region: Option<String>,
pub endpoint: Option<String>,
/// Handle for the secrets provider — never the literal credential.
pub secret_ref: Option<String>,
}
#[derive(Debug, Clone, Deserialize, Default)]
pub struct CatalogConfig {
#[serde(default = "default_manifest_prefix")]
pub manifest_prefix: String,
}
#[derive(Debug, Clone, Deserialize, Default)]
pub struct QueryConfig {
pub max_rows_per_query: Option<usize>,
}
#[derive(Debug, Clone, Deserialize)]
pub struct SidecarConfig {
#[serde(default = "default_sidecar_url")]
pub url: String,
}
#[derive(Debug, Clone, Deserialize, Default)]
pub struct AiConfig {
#[serde(default = "default_embed_model")]
pub embed_model: String,
#[serde(default = "default_gen_model")]
pub gen_model: String,
#[serde(default = "default_rerank_model")]
pub rerank_model: String,
}
#[derive(Debug, Clone, Deserialize, Default)]
pub struct AuthConfig {
#[serde(default)]
pub enabled: bool,
pub api_key: Option<String>,
}
#[derive(Debug, Clone, Deserialize, Default)]
pub struct ObservabilityConfig {
#[serde(default = "default_exporter")]
pub exporter: String,
#[serde(default = "default_service_name")]
pub service_name: String,
}
// Defaults
fn default_host() -> String { "0.0.0.0".to_string() }
fn default_gateway_port() -> u16 { 3100 }
fn default_storage_root() -> String { "./data".to_string() }
fn default_profile_root() -> String { "./data/_profiles".to_string() }
fn default_manifest_prefix() -> String { "_catalog/manifests".to_string() }
fn default_sidecar_url() -> String { "http://localhost:3200".to_string() }
fn default_embed_model() -> String { "nomic-embed-text".to_string() }
fn default_gen_model() -> String { "qwen2.5".to_string() }
fn default_rerank_model() -> String { "qwen2.5".to_string() }
fn default_exporter() -> String { "stdout".to_string() }
fn default_service_name() -> String { "lakehouse".to_string() }
impl Config {
pub fn load(path: &str) -> Result<Self, String> {
let path = Path::new(path);
if !path.exists() {
return Err(format!("config file not found: {}", path.display()));
}
let content = std::fs::read_to_string(path)
.map_err(|e| format!("failed to read config: {e}"))?;
toml::from_str(&content)
.map_err(|e| format!("failed to parse config: {e}"))
}
pub fn load_or_default() -> Self {
// Try lakehouse.toml in current dir, then /etc/lakehouse/lakehouse.toml
for path in &["lakehouse.toml", "/etc/lakehouse/lakehouse.toml"] {
if let Ok(config) = Self::load(path) {
tracing::info!("loaded config from {path}");
return config;
}
}
tracing::warn!("no config file found, using defaults");
Self::default()
}
}
impl Default for Config {
fn default() -> Self {
Self {
gateway: GatewayConfig { host: default_host(), port: default_gateway_port() },
storage: StorageConfig {
root: default_storage_root(),
profile_root: default_profile_root(),
rescue_bucket: None,
buckets: Vec::new(),
},
catalog: CatalogConfig::default(),
query: QueryConfig::default(),
sidecar: SidecarConfig { url: default_sidecar_url() },
ai: AiConfig::default(),
auth: AuthConfig::default(),
observability: ObservabilityConfig::default(),
agent: AgentSettings::default(),
}
}
}
+35
View File
@@ -0,0 +1,35 @@
use thiserror::Error;
#[derive(Debug, Error)]
pub enum LakehouseError {
#[error("dataset not found: {0}")]
DatasetNotFound(String),
#[error("object not found: {bucket}/{key}")]
ObjectNotFound { bucket: String, key: String },
#[error("storage error: {0}")]
Storage(String),
#[error("catalog error: {0}")]
Catalog(String),
#[error("query error: {0}")]
Query(String),
#[error("ai bridge error: {0}")]
AiBridge(String),
#[error("serialization error: {0}")]
Serialization(String),
}
impl LakehouseError {
pub fn status_code(&self) -> u16 {
match self {
Self::DatasetNotFound(_) | Self::ObjectNotFound { .. } => 404,
Self::Storage(_) | Self::Catalog(_) | Self::Query(_) | Self::AiBridge(_) => 500,
Self::Serialization(_) => 400,
}
}
}
+8
View File
@@ -0,0 +1,8 @@
pub mod types;
pub mod errors;
pub mod arrow_helpers;
pub mod config;
pub mod pii;
pub mod secrets;
pub mod model_matrix;
pub mod profiles;
+69
View File
@@ -0,0 +1,69 @@
//! Per-model token accounting. Entry point for the ModelMatrix work
//! the aibridge `context::estimate_tokens` deprecation has been pointing
//! at. Starts minimal — just `estimate_tokens` — so call sites can
//! migrate off the deprecated helper. Extend with per-model context
//! windows, max_tokens defaults, provider hints, etc. as we move the
//! rest of `aibridge::context::known_windows` over.
/// Namespace for per-model token + context accounting. Methods are
/// associated functions — no instance required — because the underlying
/// estimates are deterministic and stateless.
pub struct ModelMatrix;
impl ModelMatrix {
/// Rough token count — char count divided by 4, rounded up. This
/// is the same heuristic OpenAI's cookbook uses for English text;
/// it's within ±15% of BPE tokenizers for code + prose and doesn't
/// require a tokenizer lookup. Good enough for budget math where
/// the goal is "don't blow the context window" rather than exact
/// billing.
///
/// Moved from `aibridge::context::estimate_tokens` (still there with
/// a `#[deprecated]` pointer — callers should migrate here). Empty
/// string → 0; one char → 1 (ceiling of 1/4 = 1).
pub fn estimate_tokens(text: &str) -> usize {
(text.chars().count() + 3) / 4
}
}
#[cfg(test)]
mod tests {
use super::ModelMatrix;
#[test]
fn empty_string_is_zero_tokens() {
assert_eq!(ModelMatrix::estimate_tokens(""), 0);
}
#[test]
fn three_chars_is_one_token() {
// 3 → ceil(3/4) = 1. Matches the deprecated helper's behavior
// so the migration is a drop-in replacement.
assert_eq!(ModelMatrix::estimate_tokens("abc"), 1);
}
#[test]
fn four_chars_is_one_token() {
assert_eq!(ModelMatrix::estimate_tokens("abcd"), 1);
}
#[test]
fn five_chars_is_two_tokens() {
assert_eq!(ModelMatrix::estimate_tokens("abcde"), 2);
}
#[test]
fn counts_chars_not_bytes() {
// Multi-byte UTF-8 chars count as 1 char each — important for
// prompts with emoji or non-ASCII text. "héllo" is 5 chars
// (5 unicode scalars) → ceil(5/4) = 2 tokens, same as "hello".
assert_eq!(ModelMatrix::estimate_tokens("héllo"), 2);
assert_eq!(ModelMatrix::estimate_tokens("📚📚📚📚"), 1); // 4 chars
}
#[test]
fn large_text_scales_linearly() {
assert_eq!(ModelMatrix::estimate_tokens(&"x".repeat(400)), 100);
assert_eq!(ModelMatrix::estimate_tokens(&"x".repeat(401)), 101);
}
}
+116
View File
@@ -0,0 +1,116 @@
/// Auto-detect PII columns by name patterns.
/// Conservative: flags likely PII, doesn't miss obvious cases.
use crate::types::Sensitivity;
/// Check if a column name suggests PII content.
pub fn detect_sensitivity(column_name: &str) -> Option<Sensitivity> {
let lower = column_name.to_lowercase();
// Direct PII identifiers
if matches!(lower.as_str(),
"ssn" | "social_security" | "social_security_number" |
"sin" | "national_id" | "passport" | "passport_number" |
"drivers_license" | "driver_license" | "dl_number"
) {
return Some(Sensitivity::Pii);
}
// Names
if lower.contains("first_name") || lower.contains("last_name") ||
lower.contains("full_name") || lower.contains("middle_name") ||
lower == "name" || lower == "fname" || lower == "lname" {
return Some(Sensitivity::Pii);
}
// Contact info
if lower.contains("email") || lower.contains("e_mail") ||
lower.contains("phone") || lower.contains("mobile") || lower.contains("cell") ||
lower.contains("fax") || lower == "tel" || lower == "telephone" {
return Some(Sensitivity::Pii);
}
// Address
if lower.contains("address") || lower.contains("street") ||
(lower.contains("zip") && !lower.contains("unzip")) ||
lower == "postal_code" || lower == "postcode" {
return Some(Sensitivity::Pii);
}
// Financial
if lower.contains("salary") || lower.contains("wage") ||
lower.contains("pay_rate") || lower.contains("bill_rate") ||
lower.contains("compensation") || lower.contains("revenue") ||
lower.contains("bank_account") || lower.contains("routing_number") ||
lower.contains("credit_card") {
return Some(Sensitivity::Financial);
}
// Health
if lower.contains("diagnosis") || lower.contains("medication") ||
lower.contains("medical") || lower.contains("health") ||
lower.contains("patient_id") || lower.contains("mrn") {
return Some(Sensitivity::Phi);
}
// Date of birth
if lower == "dob" || lower == "date_of_birth" || lower == "birthdate" || lower == "birth_date" {
return Some(Sensitivity::Pii);
}
None
}
/// Classify all columns and return the highest sensitivity found.
pub fn detect_dataset_sensitivity(column_names: &[&str]) -> Option<Sensitivity> {
let mut highest: Option<Sensitivity> = None;
for name in column_names {
if let Some(sens) = detect_sensitivity(name) {
highest = Some(match (&highest, &sens) {
(None, s) => s.clone(),
(Some(Sensitivity::Public), s) => s.clone(),
(Some(Sensitivity::Internal), s) if !matches!(s, Sensitivity::Public | Sensitivity::Internal) => s.clone(),
(Some(Sensitivity::Financial), Sensitivity::Pii | Sensitivity::Phi) => sens.clone(),
(Some(Sensitivity::Pii), Sensitivity::Phi) => Sensitivity::Phi,
(existing, _) => existing.clone().unwrap(),
});
}
}
highest
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn detects_email_as_pii() {
assert_eq!(detect_sensitivity("email"), Some(Sensitivity::Pii));
assert_eq!(detect_sensitivity("contact_email"), Some(Sensitivity::Pii));
}
#[test]
fn detects_salary_as_financial() {
assert_eq!(detect_sensitivity("salary"), Some(Sensitivity::Financial));
assert_eq!(detect_sensitivity("bill_rate"), Some(Sensitivity::Financial));
}
#[test]
fn detects_ssn() {
assert_eq!(detect_sensitivity("ssn"), Some(Sensitivity::Pii));
}
#[test]
fn non_sensitive_returns_none() {
assert_eq!(detect_sensitivity("status"), None);
assert_eq!(detect_sensitivity("created_at"), None);
}
#[test]
fn dataset_sensitivity_picks_highest() {
let cols = vec!["id", "name", "email", "status"];
assert_eq!(detect_dataset_sensitivity(&cols), Some(Sensitivity::Pii));
}
}
+14
View File
@@ -0,0 +1,14 @@
//! ExecutionProfile — the Phase 41 rename of Phase 17's ModelProfile.
//!
//! Carries what's needed to RUN inference: model tag, dataset bindings,
//! HNSW config, embed model, bucket binding. Today this is a type
//! alias over `crate::types::ModelProfile` — the PRD's
//! "Backward compat: ModelProfile still loads, aliased to
//! ExecutionProfile" line, honored literally.
//!
//! When the migration off the old name finishes, this file can either
//! absorb the full struct definition or continue as an alias. Callers
//! should reference `ExecutionProfile` going forward; `ModelProfile`
//! stays exported from `types` for on-disk schema compat.
pub use crate::types::ModelProfile as ExecutionProfile;
+38
View File
@@ -0,0 +1,38 @@
//! MemoryProfile — how the agent's execution memory is kept.
//!
//! Phase 41 decomposition: the Phase 19 playbook_memory + Phase 26
//! successful_playbooks + Phase 45 doc_refs all need per-profile
//! tuning. Rather than bolt those onto ExecutionProfile, they live
//! here so a "thin" execution profile can reuse a "fat" memory
//! profile and vice versa.
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct MemoryProfile {
pub id: String,
#[serde(default)]
pub description: String,
/// Phase 19: ceiling for playbook_memory boost on retrieval. 0
/// disables the boost entirely.
#[serde(default = "default_boost_ceiling")]
pub playbook_boost_ceiling: f32,
/// Phase 26: max history entries retained before rotation.
#[serde(default = "default_history_cap")]
pub history_cap: usize,
/// Phase 45: stale threshold for doc_refs before drift check
/// fires (hours).
#[serde(default = "default_stale_hours")]
pub doc_stale_hours: u32,
/// Phase 28: auto-retire playbooks that fail 3+ consecutive runs.
#[serde(default = "default_true")]
pub auto_retire_on_failure: bool,
pub created_at: chrono::DateTime<chrono::Utc>,
#[serde(default)]
pub created_by: String,
}
fn default_boost_ceiling() -> f32 { 0.35 }
fn default_history_cap() -> usize { 1000 }
fn default_stale_hours() -> u32 { 168 } // one week
fn default_true() -> bool { true }
+28
View File
@@ -0,0 +1,28 @@
//! Phase 41 profile types.
//!
//! The existing `ModelProfile` (Phase 17) is aliased as
//! `ExecutionProfile` here — it continues to carry the model +
//! bindings + HNSW config needed to run inference. Three new profile
//! types land alongside: `RetrievalProfile`, `MemoryProfile`,
//! `ObserverProfile` — each owns a distinct slice of what used to be
//! bundled.
//!
//! Backward-compat rule (PRD Phase 41): existing `ModelProfile` on
//! disk continues to deserialize unchanged. New fields on the new
//! profile types are `#[serde(default)]` so old payloads load with
//! empty defaults.
//!
//! These are the canonical shapes — downstream code converts via
//! `From<ModelProfile> for ExecutionProfile` (they're the same struct
//! today, just named differently) and constructs the other three
//! as needed.
pub mod execution;
pub mod retrieval;
pub mod memory;
pub mod observer;
pub use execution::ExecutionProfile;
pub use memory::MemoryProfile;
pub use observer::ObserverProfile;
pub use retrieval::RetrievalProfile;
+38
View File
@@ -0,0 +1,38 @@
//! ObserverProfile — how loudly the observer logs this workload.
//!
//! Phase 41 decomposition: the observer's alert thresholds, escalation
//! cadence, and log retention need per-workload tuning. Hot-path
//! staffing workflows want aggressive alerting; batch backfills want
//! quieter. This profile is read by mcp-server/observer.ts at
//! activation-time.
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ObserverProfile {
pub id: String,
#[serde(default)]
pub description: String,
/// How many consecutive failures trigger a cluster escalation to
/// LLM Team `/v1/chat` (qwen3-coder:480b).
#[serde(default = "default_failure_cluster_size")]
pub failure_cluster_size: u32,
/// Minimum seconds between alert emails for the same sig_hash.
/// Prevents alert storms during a regression.
#[serde(default = "default_alert_cooldown")]
pub alert_cooldown_secs: u32,
/// Observer ring buffer size. Older events fall off when full.
#[serde(default = "default_ring_size")]
pub ring_size: usize,
/// Whether to forward events to external Langfuse
/// (/v1/langfuse_trace). Off by default.
#[serde(default)]
pub forward_to_langfuse: bool,
pub created_at: chrono::DateTime<chrono::Utc>,
#[serde(default)]
pub created_by: String,
}
fn default_failure_cluster_size() -> u32 { 3 }
fn default_alert_cooldown() -> u32 { 300 } // 5 minutes
fn default_ring_size() -> usize { 2000 }
+52
View File
@@ -0,0 +1,52 @@
//! RetrievalProfile — what + how the agent reaches into memory.
//!
//! Phase 41 decomposition: the old ModelProfile bundled "what dataset
//! can I read" (bound_datasets) AND "how do I rank results"
//! (hnsw_config) with the model tag. Retrieval concerns split out here
//! so a profile can swap its retrieval strategy without re-activating
//! the model.
//!
//! Fields chosen for what's actually varied per-workload today:
//! - `top_k` / `rerank_top_k` — how many hits to fetch + rerank
//! - `freshness_cutoff_days` — Phase 45 doc-drift uses this
//! - `boost_playbook_memory` — Phase 19 meta-index feedback
//! - `enforce_sensitivity_gates` — Phase 13 access-control integration
//!
//! All fields are `#[serde(default)]` so loading a profile file that
//! predates Phase 41 works without migration.
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct RetrievalProfile {
/// Unique id — slug form, separate namespace from ExecutionProfile.
pub id: String,
/// Free-text operator description.
#[serde(default)]
pub description: String,
/// Default top-K for /vectors/search + /vectors/hybrid.
#[serde(default = "default_top_k")]
pub top_k: u32,
/// How many of the top-K to pass through the reranker. 0 disables
/// reranking for this profile.
#[serde(default = "default_rerank_top_k")]
pub rerank_top_k: u32,
/// Don't consider playbooks / docs older than this (days). 0 or
/// absent = no freshness filter.
#[serde(default)]
pub freshness_cutoff_days: u32,
/// Phase 19: boost workers/results by playbook_memory similarity.
#[serde(default)]
pub boost_playbook_memory: bool,
/// Phase 13: apply access-control masking on sensitive columns.
/// Default on — safety-first.
#[serde(default = "default_true")]
pub enforce_sensitivity_gates: bool,
pub created_at: chrono::DateTime<chrono::Utc>,
#[serde(default)]
pub created_by: String,
}
fn default_top_k() -> u32 { 10 }
fn default_rerank_top_k() -> u32 { 5 }
fn default_true() -> bool { true }
+144
View File
@@ -0,0 +1,144 @@
/// Secrets layer for bucket credentials.
///
/// Credentials never appear in `lakehouse.toml`. Each bucket entry has an
/// opaque `secret_ref` handle, resolved at startup by a `SecretsProvider`.
///
/// MVP ships `FileSecretsProvider` — reads a TOML file at a path given by
/// `LAKEHOUSE_SECRETS` env var, defaulting to `/etc/lakehouse/secrets.toml`.
/// The file should be root-owned, mode 0600, never in git.
///
/// Future providers (Vault, SOPS, OS keyring) implement the same trait.
use async_trait::async_trait;
use serde::Deserialize;
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
#[derive(Debug, Clone)]
pub struct BucketCredentials {
pub access_key: String,
pub secret_key: String,
}
// Redact credentials when debug-printing — NEVER let these end up in logs.
impl BucketCredentials {
pub fn redacted(&self) -> String {
let prefix: String = self.access_key.chars().take(4).collect();
format!("access_key={}... (redacted)", prefix)
}
}
#[async_trait]
pub trait SecretsProvider: Send + Sync {
/// Resolve a handle to credentials. Returns Err if the handle is unknown.
async fn resolve(&self, handle: &str) -> Result<BucketCredentials, String>;
/// List known handles (for diagnostics only, never values).
async fn list_handles(&self) -> Vec<String>;
}
#[derive(Debug, Deserialize)]
struct RawSecretEntry {
access_key: String,
secret_key: String,
}
/// File-backed secrets provider.
pub struct FileSecretsProvider {
path: PathBuf,
cache: tokio::sync::RwLock<HashMap<String, BucketCredentials>>,
}
impl FileSecretsProvider {
pub fn new(path: impl Into<PathBuf>) -> Self {
Self {
path: path.into(),
cache: tokio::sync::RwLock::new(HashMap::new()),
}
}
/// Default path: env var `LAKEHOUSE_SECRETS` or `/etc/lakehouse/secrets.toml`.
pub fn from_env() -> Self {
let path = std::env::var("LAKEHOUSE_SECRETS")
.unwrap_or_else(|_| "/etc/lakehouse/secrets.toml".to_string());
Self::new(path)
}
/// Arc wrapper for convenient injection.
pub fn shared(path: impl Into<PathBuf>) -> Arc<dyn SecretsProvider> {
Arc::new(Self::new(path))
}
/// Load and populate the cache. Call once at startup.
pub async fn load(&self) -> Result<usize, String> {
if !self.path.exists() {
tracing::info!(
"secrets file {} not present — running with zero credentials configured",
self.path.display()
);
return Ok(0);
}
check_permissions(&self.path)?;
let content = std::fs::read_to_string(&self.path)
.map_err(|e| format!("read secrets file: {e}"))?;
let raw: HashMap<String, RawSecretEntry> = toml::from_str(&content)
.map_err(|e| format!("parse secrets file: {e}"))?;
let mut cache = self.cache.write().await;
for (handle, entry) in raw {
cache.insert(handle.clone(), BucketCredentials {
access_key: entry.access_key,
secret_key: entry.secret_key,
});
}
let count = cache.len();
tracing::info!("secrets: loaded {count} handles from {}", self.path.display());
Ok(count)
}
}
#[async_trait]
impl SecretsProvider for FileSecretsProvider {
async fn resolve(&self, handle: &str) -> Result<BucketCredentials, String> {
let cache = self.cache.read().await;
cache
.get(handle)
.cloned()
.ok_or_else(|| format!("secret handle not found: {handle}"))
}
async fn list_handles(&self) -> Vec<String> {
self.cache.read().await.keys().cloned().collect()
}
}
/// Reject world-readable secrets files. Fail closed.
fn check_permissions(path: &Path) -> Result<(), String> {
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let meta = std::fs::metadata(path)
.map_err(|e| format!("stat secrets file: {e}"))?;
let mode = meta.permissions().mode() & 0o777;
if mode & 0o044 != 0 {
return Err(format!(
"secrets file {} is group/world-readable (mode {:o}); chmod 600 to proceed",
path.display(), mode
));
}
}
let _ = path;
Ok(())
}
/// A zero-handles provider for tests / purely-local setups.
pub struct EmptySecretsProvider;
#[async_trait]
impl SecretsProvider for EmptySecretsProvider {
async fn resolve(&self, handle: &str) -> Result<BucketCredentials, String> {
Err(format!("no secrets provider configured (wanted handle '{handle}')"))
}
async fn list_handles(&self) -> Vec<String> { Vec::new() }
}
+366
View File
@@ -0,0 +1,366 @@
use serde::{Deserialize, Serialize};
use uuid::Uuid;
/// Unique identifier for a dataset in the catalog.
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct DatasetId(pub Uuid);
impl DatasetId {
pub fn new() -> Self {
Self(Uuid::new_v4())
}
}
impl std::fmt::Display for DatasetId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
/// Pointer to an object in storage. This is what the catalog stores — never raw data.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ObjectRef {
pub bucket: String,
pub key: String,
pub size_bytes: u64,
pub created_at: chrono::DateTime<chrono::Utc>,
}
/// Schema fingerprint — deterministic hash of an Arrow schema.
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct SchemaFingerprint(pub String);
/// Sensitivity classification for a field or dataset.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Sensitivity {
Public, // anyone can see
Internal, // employees only
Pii, // personally identifiable (name, email, phone, SSN)
Phi, // protected health info (HIPAA)
Financial, // compensation, billing, revenue
}
/// Column-level metadata.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ColumnMeta {
pub name: String,
pub data_type: String,
#[serde(default)]
pub sensitivity: Option<Sensitivity>,
#[serde(default)]
pub description: String,
#[serde(default)]
pub is_pii: bool,
}
/// Data lineage — where this dataset came from.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Lineage {
pub source_system: String, // "csv_upload", "mysql_import", "api_ingest"
pub source_file: String, // original filename or URI
pub ingest_job: String, // job ID that created this dataset
pub ingest_timestamp: chrono::DateTime<chrono::Utc>,
#[serde(default)]
pub parent_datasets: Vec<String>, // derived from these datasets
}
/// Freshness contract — how often this data should be updated.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FreshnessContract {
pub expected_frequency: String, // "daily", "weekly", "hourly", "manual"
pub stale_after_hours: u32, // alert if not updated within this many hours
}
/// Dataset manifest — the catalog's complete view of a dataset.
/// All new fields use #[serde(default)] for backward compatibility with existing manifests.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DatasetManifest {
pub id: DatasetId,
pub name: String,
pub schema_fingerprint: SchemaFingerprint,
pub objects: Vec<ObjectRef>,
pub created_at: chrono::DateTime<chrono::Utc>,
pub updated_at: chrono::DateTime<chrono::Utc>,
// --- Rich metadata (Phase 10) ---
/// Human-readable description (auto-generated on ingest or manual)
#[serde(default)]
pub description: String,
/// Who owns this dataset
#[serde(default)]
pub owner: String,
/// Overall sensitivity classification
#[serde(default)]
pub sensitivity: Option<Sensitivity>,
/// Column-level metadata including sensitivity and descriptions
#[serde(default)]
pub columns: Vec<ColumnMeta>,
/// Where this data came from
#[serde(default)]
pub lineage: Option<Lineage>,
/// How fresh the data should be
#[serde(default)]
pub freshness: Option<FreshnessContract>,
/// Free-form tags for discovery
#[serde(default)]
pub tags: Vec<String>,
/// Row count (updated on ingest/compact)
#[serde(default)]
pub row_count: Option<u64>,
// --- Embedding freshness (Phase C) ---
/// When the attached vector index was last refreshed. `None` means this
/// dataset has no vector index yet, or has never been embedded.
#[serde(default)]
pub last_embedded_at: Option<chrono::DateTime<chrono::Utc>>,
/// When data was written that hasn't yet been embedded. `Some(t)` means
/// the vector index is out-of-date as of timestamp `t`. Cleared on refresh.
#[serde(default)]
pub embedding_stale_since: Option<chrono::DateTime<chrono::Utc>>,
/// How this dataset wants stale embeddings handled.
#[serde(default)]
pub embedding_refresh_policy: Option<RefreshPolicy>,
}
/// Controls what happens when new data is written to a dataset with an
/// attached vector index.
///
/// - `Manual` (default): data writes set `embedding_stale_since`; nothing
/// embeds until an operator or agent calls `/vectors/refresh/{dataset}`.
/// - `OnAppend`: ingest fires a background refresh immediately after writing.
/// Suitable for datasets where vector freshness matters more than ingest
/// latency.
/// - `Scheduled(cron)`: a timer or external scheduler triggers refresh at
/// the named cadence. The scheduler itself is not in this ADR scope —
/// this just declares the intent so operators can see which policy a
/// dataset expects.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum RefreshPolicy {
Manual,
OnAppend,
Scheduled { cron: String },
}
impl Default for RefreshPolicy {
fn default() -> Self { Self::Manual }
}
/// AI-safe view of a dataset (Phase D).
///
/// A view is a named projection over a base dataset that:
/// - whitelists the columns visible (everything else is invisible),
/// - optionally applies a row filter (a SQL WHERE clause fragment),
/// - optionally redacts column values (e.g. mask the last 4 of an SSN).
///
/// Queries that reference the view name see only the safe projection.
/// The base table remains separately accessible by its raw name — this
/// is a *capability surface*, not a data classification fence. Use
/// access control (Phase 13) on top to enforce who can query which view.
///
/// Stored as JSON at `_catalog/views/{name}.json` in the primary bucket
/// alongside dataset manifests. The view itself doesn't store data —
/// queryd materializes it as a DataFusion view at SessionContext build
/// time.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AiView {
pub name: String,
/// Source dataset name. Must exist in the catalog.
pub base_dataset: String,
/// Whitelisted columns. Empty list is invalid (would expose nothing).
pub columns: Vec<String>,
/// Optional WHERE clause fragment (without the `WHERE` keyword), e.g.
/// `status != 'blocked' AND city = 'Chicago'`. SQL injection here is
/// the operator's problem — only trust this from authenticated admins.
#[serde(default)]
pub row_filter: Option<String>,
/// Per-column redactions: `{column_name -> Redaction}`.
#[serde(default)]
pub column_redactions: std::collections::HashMap<String, Redaction>,
pub created_at: chrono::DateTime<chrono::Utc>,
#[serde(default)]
pub created_by: String,
#[serde(default)]
pub description: String,
}
/// Per-model configuration for HNSW builds. Declared inline so the PRD
/// invariant 9 ("every reader gets its own profile") can carry tuning
/// preferences without forcing a round-trip to vectord's trial system.
/// Mirrors `vectord::trial::HnswConfig` to avoid a cross-crate dep cycle.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ProfileHnswConfig {
pub ef_construction: usize,
pub ef_search: usize,
#[serde(default)]
pub seed: Option<u64>,
}
impl Default for ProfileHnswConfig {
/// Matches the production HNSW default from Phase 15 trial sweep
/// (see docs/ADR and vectord::trial::HnswConfig::default).
fn default() -> Self {
Self { ef_construction: 80, ef_search: 30, seed: None }
}
}
/// A named client of the substrate — a local LLM, an AI agent, or a
/// human role that should see a specific scoped view of the data.
/// Profiles are the first-class "every reader gets its own profile"
/// concept from PRD invariant 9.
///
/// Bound datasets can be raw catalog tables OR `AiView` names — either
/// works at query time since both register as DataFusion tables with
/// the same name. Mixing raw and views lets you compose access rules:
/// bind a model to `candidates_safe` (PII-redacted view) and `placements`
/// (raw table, no PII), and the model gets exactly what it needs.
///
/// Stored at `_catalog/profiles/{id}.json`. Activation (`/profile/{id}/
/// activate`) pre-loads the embedding cache and builds HNSW with the
/// profile's config — first search after activate is warm.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelProfile {
/// Unique id — slug form (alphanumeric + `-`, `_`). Used in URLs,
/// bucket names (`profile:{id}`), and trial journal prefixes.
pub id: String,
/// Ollama model tag (e.g. "qwen2.5:7b", "nomic-embed-text:latest").
/// Used when this profile is the "generate" or "embed" caller.
pub ollama_name: String,
/// Free-text description surfaced to operators. Who/what is this
/// profile for?
#[serde(default)]
pub description: String,
/// Names of datasets or AiViews this profile can read. A search
/// against an index whose `source` isn't in this list is refused.
pub bound_datasets: Vec<String>,
/// Per-profile HNSW config, used when activate builds the graph.
#[serde(default)]
pub hnsw_config: ProfileHnswConfig,
/// Which embedding model this profile expects. Defaults to
/// `nomic-embed-text` (the only one currently installed).
#[serde(default = "default_embed_model")]
pub embed_model: String,
pub created_at: chrono::DateTime<chrono::Utc>,
#[serde(default)]
pub created_by: String,
/// Federation layer 2: which bucket this profile's artifacts
/// (trial journals, promotions for its indexes) live in. If set and
/// not yet registered, `POST /vectors/profile/{id}/activate` auto-
/// provisions a local bucket under `storage.profile_root` and
/// registers it before warming indexes. `None` = share the primary
/// bucket (pre-federation behavior).
#[serde(default)]
pub bucket: Option<String>,
/// ADR-019 hybrid: which storage backend serves this profile's
/// indexes. Default Parquet so pre-existing profiles unchanged.
/// New profiles can opt into Lance for append-heavy / random-access
/// / 5M+ vector workloads.
#[serde(default)]
pub vector_backend: VectorBackend,
/// Phase 41: Profile type for routing + activation behavior.
/// Default Execution for backward compatibility.
#[serde(default)]
pub profile_type: ProfileType,
}
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum ProfileType {
#[default]
Execution,
Retrieval,
Memory,
Observer,
}
fn default_embed_model() -> String { "nomic-embed-text".to_string() }
/// Per-profile / per-index storage backend choice.
///
/// `Parquet` (default) — vectors as binary-blob in Parquet + in-RAM HNSW
/// for ANN. Best when:
/// - Index ≤ ~5M vectors (fits in RAM)
/// - Search latency dominates (HNSW p50 < 1ms at 100K)
/// - Workload is mostly read, not append
///
/// `Lance` — vectors as native FixedSizeList in a Lance dataset, IVF_PQ
/// on-disk index. Best when:
/// - Index will exceed RAM ceiling (5M+ vectors per index)
/// - Heavy random-row access for RAG (Lance: 311us; Parquet: ~35ms full
/// scan per fetch)
/// - Append-heavy ingest (Lance: 0.08s for 10K rows; Parquet: full
/// ~330MB rewrite per change)
/// - Need to rebuild index frequently (Lance IVF_PQ: 16s; HNSW: 230s
/// at the same scale)
///
/// Defaults to Parquet for backward compat with pre-ADR-019 indexes.
/// See `docs/ADR-019-vector-storage.md` for the full scorecard.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum VectorBackend {
Parquet,
Lance,
}
impl Default for VectorBackend {
fn default() -> Self { VectorBackend::Parquet }
}
/// Soft-delete marker (Phase E).
///
/// Tombstones live beside the dataset in `_catalog/tombstones/{dataset}/`
/// as append-log JSONL. Query-time filter in queryd reads all tombstones
/// for each dataset and wraps the base table in a DataFusion view that
/// excludes tombstoned rows. Physical deletion happens later (compaction),
/// so the row count immediately reflects the delete but data is still on
/// disk until compact runs. That's deliberate — it gives a reversal
/// window and keeps the event journal audit trail intact.
///
/// All tombstones for a given dataset must use the same `row_key_column`
/// (enforced on write); otherwise the query-time filter can't be built
/// as a single WHERE clause.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Tombstone {
pub dataset: String,
/// Column name that identifies the row (e.g. "candidate_id").
pub row_key_column: String,
/// Value of that column for the tombstoned row.
pub row_key_value: String,
pub deleted_at: chrono::DateTime<chrono::Utc>,
/// Human / system actor responsible for the delete (audit).
#[serde(default)]
pub actor: String,
/// Why (e.g. "GDPR request #1234", "user-requested erasure").
#[serde(default)]
pub reason: String,
}
/// How a column's values should be transformed before being returned.
/// `Mask` is the most common — keeps a few visible chars, replaces the
/// rest with `*`. `Hash` returns SHA-256 of the value for join keys you
/// want to preserve referential integrity on without exposing the value.
/// `Null` returns NULL unconditionally.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum Redaction {
Null,
Hash,
Mask {
/// Number of leading characters left visible.
keep_prefix: usize,
/// Number of trailing characters left visible.
keep_suffix: usize,
},
}
+16
View File
@@ -0,0 +1,16 @@
[package]
name = "storaged"
version = "0.1.0"
edition = "2024"
[dependencies]
shared = { path = "../shared" }
tokio = { workspace = true }
axum = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }
tracing = { workspace = true }
object_store = { workspace = true }
bytes = { workspace = true }
futures = { workspace = true }
chrono = { workspace = true }
+330
View File
@@ -0,0 +1,330 @@
/// Write-once batched append log.
///
/// Problem we're fixing: the error journal and HNSW trial journal both
/// previously did read-modify-write of their whole JSONL file on every
/// event. That's O(N²) cumulative work and generates huge churn at scale.
/// It's exactly the pattern llms3.com flags as the "small-file /
/// rewrite-amplification" anti-pattern.
///
/// This helper implements the pattern object storage actually wants:
///
/// - Events accumulate in an **in-memory buffer** (reads see them immediately).
/// - When the buffer hits a threshold, or `flush()` is called, the buffer is
/// written **as one new object** with a timestamp-sorted key.
/// - Existing objects are never rewritten.
/// - Reads enumerate all batch files, sort by key, and concat in order.
/// - An explicit `compact()` reads every batch file, writes one consolidated
/// file, and deletes the originals — the LSM-tree compaction idea applied
/// to small JSONL events.
///
/// Storage layout:
/// ```text
/// {prefix}/
/// batch_0001776319628000123.jsonl
/// batch_0001776319745987654.jsonl
/// batch_compacted_00001776319800000000.jsonl (after compact())
/// ```
/// Key format: the zero-padded epoch microsecond of the write, so lexical
/// sort == chronological sort.
use bytes::Bytes;
use chrono::Utc;
use object_store::ObjectStore;
use std::sync::Arc;
use tokio::sync::Mutex;
use crate::ops;
const DEFAULT_FLUSH_THRESHOLD: usize = 32;
pub struct AppendLog {
store: Arc<dyn ObjectStore>,
prefix: String,
buffer: Mutex<Vec<Vec<u8>>>,
flush_threshold: usize,
}
impl AppendLog {
/// Create a new append log rooted at `prefix` in the given object store.
/// Events auto-flush when buffer reaches `flush_threshold` (default 32).
pub fn new(store: Arc<dyn ObjectStore>, prefix: impl Into<String>) -> Self {
Self {
store,
prefix: prefix.into(),
buffer: Mutex::new(Vec::new()),
flush_threshold: DEFAULT_FLUSH_THRESHOLD,
}
}
pub fn with_flush_threshold(mut self, threshold: usize) -> Self {
self.flush_threshold = threshold.max(1);
self
}
/// Add an event. The returned future completes either immediately
/// (buffered) or after a flush, depending on whether the buffer hit the
/// threshold. Callers don't need to care either way.
pub async fn append(&self, line: Vec<u8>) -> Result<(), String> {
let should_flush = {
let mut buf = self.buffer.lock().await;
buf.push(line);
buf.len() >= self.flush_threshold
};
if should_flush {
self.flush().await?;
}
Ok(())
}
/// Force-flush the in-memory buffer to object storage as a single new
/// batch file. Safe to call anytime (idempotent no-op when buffer empty).
pub async fn flush(&self) -> Result<(), String> {
let batch = {
let mut buf = self.buffer.lock().await;
if buf.is_empty() { return Ok(()); }
std::mem::take(&mut *buf)
};
let ts_us = Utc::now().timestamp_micros().max(0) as u128;
let key = format!("{}/batch_{:019}.jsonl", self.prefix, ts_us);
let mut body = Vec::with_capacity(batch.iter().map(|b| b.len() + 1).sum());
for line in batch {
body.extend_from_slice(&line);
if !body.ends_with(b"\n") { body.push(b'\n'); }
}
ops::put(&self.store, &key, Bytes::from(body)).await
}
/// Read every event across all batch files + unflushed in-memory buffer.
/// Events are returned in chronological order.
pub async fn read_all(&self) -> Result<Vec<Vec<u8>>, String> {
let mut keys = self.list_batch_keys().await?;
keys.sort();
let mut out = Vec::new();
for key in keys {
let bytes = ops::get(&self.store, &key).await?;
for line in bytes.split(|b| *b == b'\n') {
if !line.is_empty() {
out.push(line.to_vec());
}
}
}
// Include unflushed events so callers see the latest state
// whether or not someone ran flush() recently.
let buf = self.buffer.lock().await;
for line in buf.iter() {
out.push(line.clone());
}
Ok(out)
}
/// Consolidate all current batch files into one compacted file, then
/// delete the originals. Safe to call while appends are in flight:
/// new batches written during compaction get a higher timestamp and
/// survive. Fails closed — if anything goes wrong mid-delete, the
/// compacted file coexists with originals and next read sees duplicates
/// (which the dedup caller must handle) rather than data loss.
pub async fn compact(&self) -> Result<CompactStats, String> {
// Snapshot which files to compact BEFORE we write the new one.
let mut originals = self.list_batch_keys().await?;
originals.sort();
if originals.len() < 2 {
return Ok(CompactStats { merged_files: originals.len(), events: 0, new_key: None });
}
// Gather all existing events.
let mut events = Vec::new();
for key in &originals {
let bytes = ops::get(&self.store, key).await?;
for line in bytes.split(|b| *b == b'\n') {
if !line.is_empty() {
events.push(line.to_vec());
}
}
}
let total_events = events.len();
if total_events == 0 {
// Clean up empty files without writing a new one.
for key in &originals {
let _ = ops::delete(&self.store, key).await;
}
return Ok(CompactStats { merged_files: originals.len(), events: 0, new_key: None });
}
// Name the compacted file with the SAME `batch_{ts}` format so it
// sorts chronologically with future batches. Using a distinct prefix
// ("batch_compacted_") would break lex ordering: later `batch_N`
// files would sort BEFORE the compacted file because 'c' > digits.
// Timestamp = now, so any appends arriving during compaction (which
// get the current wall-clock time) sort AFTER this file.
let ts_us = Utc::now().timestamp_micros().max(0) as u128;
let new_key = format!("{}/batch_{:019}.jsonl", self.prefix, ts_us);
let mut body = Vec::new();
for line in &events {
body.extend_from_slice(line);
body.push(b'\n');
}
ops::put(&self.store, &new_key, Bytes::from(body)).await?;
// Only delete originals once the consolidated file is persisted.
let mut failures = 0;
for key in &originals {
if ops::delete(&self.store, key).await.is_err() {
failures += 1;
}
}
if failures > 0 {
tracing::warn!(
"compact '{}': {} original files failed to delete — consolidated file {} has the data",
self.prefix, failures, new_key,
);
}
Ok(CompactStats {
merged_files: originals.len(),
events: total_events,
new_key: Some(new_key),
})
}
/// How many batch files exist on disk right now.
pub async fn file_count(&self) -> Result<usize, String> {
Ok(self.list_batch_keys().await?.len())
}
async fn list_batch_keys(&self) -> Result<Vec<String>, String> {
let prefix_with_slash = format!("{}/", self.prefix);
// list the prefix then filter for keys that match our naming scheme;
// unrelated files at the same prefix won't be touched.
let raw = ops::list(&self.store, Some(&prefix_with_slash)).await?;
Ok(raw
.into_iter()
.filter(|k| {
let basename = k.rsplit('/').next().unwrap_or(k);
basename.starts_with("batch_") && basename.ends_with(".jsonl")
})
.collect())
}
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct CompactStats {
pub merged_files: usize,
pub events: usize,
pub new_key: Option<String>,
}
// Log unflushed-buffer size on drop. We can't `.await` from a sync `Drop`,
// so a real flush isn't possible here — callers are responsible for calling
// `.flush()` before dropping if durability matters. These journals are
// observability hints; a few lost buffered events at shutdown are
// acceptable per ADR-018.
impl Drop for AppendLog {
fn drop(&mut self) {
let buf_len = self.buffer.try_lock().map(|b| b.len()).unwrap_or(0);
if buf_len == 0 { return; }
tracing::debug!(
"append_log '{}' dropping with {} unflushed events",
self.prefix, buf_len,
);
}
}
#[cfg(test)]
mod tests {
use super::*;
use object_store::memory::InMemory;
fn mk(threshold: usize) -> AppendLog {
AppendLog::new(Arc::new(InMemory::new()), "prefix")
.with_flush_threshold(threshold)
}
#[tokio::test]
async fn append_stays_buffered_below_threshold() {
let log = mk(5);
log.append(b"one".to_vec()).await.unwrap();
log.append(b"two".to_vec()).await.unwrap();
assert_eq!(log.file_count().await.unwrap(), 0, "no files until threshold");
let all = log.read_all().await.unwrap();
assert_eq!(all, vec![b"one".to_vec(), b"two".to_vec()],
"read_all surfaces unflushed buffer");
}
#[tokio::test]
async fn append_auto_flushes_on_threshold() {
let log = mk(3);
for i in 0..3 {
log.append(format!("evt{i}").into_bytes()).await.unwrap();
}
assert_eq!(log.file_count().await.unwrap(), 1, "threshold triggered one flush");
// A fourth append stays buffered until the next threshold.
log.append(b"evt3".to_vec()).await.unwrap();
assert_eq!(log.file_count().await.unwrap(), 1, "below threshold again");
}
#[tokio::test]
async fn flush_empty_is_noop() {
let log = mk(32);
log.flush().await.unwrap();
log.flush().await.unwrap();
assert_eq!(log.file_count().await.unwrap(), 0);
}
#[tokio::test]
async fn read_all_orders_events_across_flushes() {
let log = mk(1); // flush-on-every-append
for i in 0..5 {
log.append(format!("e{i}").into_bytes()).await.unwrap();
// Spread writes out so timestamps sort strictly.
tokio::time::sleep(std::time::Duration::from_millis(2)).await;
}
let all = log.read_all().await.unwrap();
let strs: Vec<String> = all.into_iter()
.map(|v| String::from_utf8(v).unwrap())
.collect();
assert_eq!(strs, vec!["e0", "e1", "e2", "e3", "e4"],
"lex sort of batch keys == chronological event order");
}
#[tokio::test]
async fn compact_merges_multiple_files_into_one() {
let log = mk(1); // force file-per-append
for i in 0..4 {
log.append(format!("e{i}").into_bytes()).await.unwrap();
tokio::time::sleep(std::time::Duration::from_millis(2)).await;
}
assert_eq!(log.file_count().await.unwrap(), 4);
let stats = log.compact().await.unwrap();
assert_eq!(stats.merged_files, 4);
assert_eq!(stats.events, 4);
assert!(stats.new_key.is_some());
assert_eq!(log.file_count().await.unwrap(), 1, "originals deleted, 1 survivor");
let all = log.read_all().await.unwrap();
assert_eq!(all.len(), 4, "no events lost in compaction");
}
#[tokio::test]
async fn compact_with_single_file_is_noop() {
let log = mk(1);
log.append(b"only".to_vec()).await.unwrap();
assert_eq!(log.file_count().await.unwrap(), 1);
let stats = log.compact().await.unwrap();
assert_eq!(stats.merged_files, 1);
assert_eq!(stats.events, 0, "nothing to consolidate");
assert!(stats.new_key.is_none(), "no new file written");
assert_eq!(log.file_count().await.unwrap(), 1, "original untouched");
}
}
+10
View File
@@ -0,0 +1,10 @@
use object_store::ObjectStore;
use object_store::local::LocalFileSystem;
use std::sync::Arc;
/// Initialize the object store backend.
/// Starts with LocalFileSystem. Swap to S3/RustFS by changing this function.
pub fn init_local(root: &str) -> Arc<dyn ObjectStore> {
std::fs::create_dir_all(root).expect("failed to create storage root");
Arc::new(LocalFileSystem::new_with_prefix(root).expect("failed to init local object store"))
}
+289
View File
@@ -0,0 +1,289 @@
/// Bucket operation error journal.
///
/// Every bucket op failure and every rescue fallback lands here. Goal:
/// answering "is anything broken?" with one HTTP call.
///
/// Storage: batched write-once files under `_errors/bucket_errors/` in the
/// primary bucket. Uses the shared `AppendLog` helper so we never rewrite
/// existing files — see `append_log.rs` for the full pattern rationale.
///
/// In-memory ring buffer holds the last N events for fast response; on
/// startup `load_recent` hydrates it from all batch files.
use chrono::{DateTime, Utc};
use object_store::ObjectStore;
use serde::{Deserialize, Serialize};
use std::collections::VecDeque;
use std::sync::Arc;
use tokio::sync::RwLock;
use crate::append_log::{AppendLog, CompactStats};
const JOURNAL_PREFIX: &str = "_errors/bucket_errors";
const RING_CAPACITY: usize = 2000;
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum BucketOp {
Read,
Write,
Delete,
List,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BucketErrorEvent {
pub ts: DateTime<Utc>,
pub op: BucketOp,
pub target: String,
pub key: String,
pub error: String,
#[serde(default)]
pub rescued: bool,
}
impl BucketErrorEvent {
pub fn new_read(target: &str, key: &str, error: &str) -> Self {
Self { ts: Utc::now(), op: BucketOp::Read, target: target.into(), key: key.into(), error: error.into(), rescued: false }
}
pub fn new_write(target: &str, key: &str, error: &str) -> Self {
Self { ts: Utc::now(), op: BucketOp::Write, target: target.into(), key: key.into(), error: error.into(), rescued: false }
}
pub fn new_delete(target: &str, key: &str, error: &str) -> Self {
Self { ts: Utc::now(), op: BucketOp::Delete, target: target.into(), key: key.into(), error: error.into(), rescued: false }
}
pub fn new_list(target: &str, prefix: &str, error: &str) -> Self {
Self { ts: Utc::now(), op: BucketOp::List, target: target.into(), key: prefix.into(), error: error.into(), rescued: false }
}
}
#[derive(Clone)]
pub struct ErrorJournal {
log: Arc<AppendLog>,
ring: Arc<RwLock<VecDeque<BucketErrorEvent>>>,
}
#[derive(Debug, Clone, Serialize)]
pub struct HealthReport {
pub period_minutes: i64,
pub total_errors: usize,
pub per_bucket: std::collections::HashMap<String, usize>,
pub unhealthy_buckets: Vec<String>,
}
impl ErrorJournal {
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
// Keep flush threshold lowish — operators checking /storage/errors
// after a recent incident want to see fresh rows even without an
// explicit flush.
let log = Arc::new(
AppendLog::new(store, JOURNAL_PREFIX).with_flush_threshold(8),
);
Self {
log,
ring: Arc::new(RwLock::new(VecDeque::with_capacity(RING_CAPACITY))),
}
}
/// Hydrate the ring buffer from existing batch files. Call once at
/// startup. Tolerates malformed lines (skipped with a warning) and
/// missing files (returns 0).
pub async fn load_recent(&self) -> Result<usize, String> {
let lines = self.log.read_all().await.unwrap_or_default();
let mut ring = self.ring.write().await;
for line in lines {
match serde_json::from_slice::<BucketErrorEvent>(&line) {
Ok(ev) => {
if ring.len() >= RING_CAPACITY { ring.pop_front(); }
ring.push_back(ev);
}
Err(e) => tracing::warn!("error journal: skip malformed line ({e})"),
}
}
Ok(ring.len())
}
/// Append an event. In-memory ring updated immediately; persistence
/// happens in batches via the underlying AppendLog.
pub async fn append(&self, event: BucketErrorEvent) {
{
let mut ring = self.ring.write().await;
if ring.len() >= RING_CAPACITY { ring.pop_front(); }
ring.push_back(event.clone());
}
match serde_json::to_vec(&event) {
Ok(line) => {
if let Err(e) = self.log.append(line).await {
tracing::error!("error journal persist failed: {e}");
}
}
Err(e) => tracing::error!("error journal serialize failed: {e}"),
}
}
/// Mark the most recent matching in-memory event as rescued.
/// Only updates the ring buffer — the JSONL line already persisted
/// records the failure fact; rescue status travels in response headers.
pub async fn mark_rescued_last(&self, target: &str, key: &str) {
let mut ring = self.ring.write().await;
for ev in ring.iter_mut().rev() {
if matches!(ev.op, BucketOp::Read) && ev.target == target && ev.key == key && !ev.rescued {
ev.rescued = true;
break;
}
}
}
pub async fn recent(&self, limit: usize) -> Vec<BucketErrorEvent> {
let ring = self.ring.read().await;
let start = ring.len().saturating_sub(limit);
ring.iter().skip(start).cloned().collect()
}
pub async fn filter(
&self,
bucket: Option<&str>,
since: Option<DateTime<Utc>>,
limit: usize,
) -> Vec<BucketErrorEvent> {
let ring = self.ring.read().await;
ring.iter()
.rev()
.filter(|ev| bucket.map_or(true, |b| ev.target == b))
.filter(|ev| since.map_or(true, |s| ev.ts >= s))
.take(limit)
.cloned()
.collect::<Vec<_>>()
.into_iter()
.rev()
.collect()
}
/// Summarize errors in the last N minutes.
pub async fn health(&self, period_minutes: i64) -> HealthReport {
use std::collections::HashMap;
let cutoff = Utc::now() - chrono::Duration::minutes(period_minutes);
let ring = self.ring.read().await;
let recent: Vec<_> = ring.iter().filter(|ev| ev.ts >= cutoff).collect();
let mut per_bucket: HashMap<String, usize> = HashMap::new();
for ev in &recent {
*per_bucket.entry(ev.target.clone()).or_insert(0) += 1;
}
let unhealthy_buckets: Vec<String> = per_bucket
.iter()
.filter(|(_, c)| **c >= 3)
.map(|(k, _)| k.clone())
.collect();
HealthReport {
period_minutes,
total_errors: recent.len(),
per_bucket,
unhealthy_buckets,
}
}
/// Force an immediate flush of buffered events to object storage.
pub async fn flush(&self) -> Result<(), String> {
self.log.flush().await
}
/// Consolidate all batch files into one. Operator cleanup.
pub async fn compact(&self) -> Result<CompactStats, String> {
self.log.compact().await
}
/// How many JSONL batch files currently exist.
pub async fn file_count(&self) -> Result<usize, String> {
self.log.file_count().await
}
}
#[cfg(test)]
mod tests {
use super::*;
use object_store::memory::InMemory;
fn mk() -> ErrorJournal {
ErrorJournal::new(Arc::new(InMemory::new()))
}
#[tokio::test]
async fn append_updates_ring_immediately() {
let j = mk();
j.append(BucketErrorEvent::new_read("tenant_a", "k1", "timeout")).await;
j.append(BucketErrorEvent::new_write("tenant_b", "k2", "denied")).await;
let recent = j.recent(10).await;
assert_eq!(recent.len(), 2);
assert_eq!(recent[0].target, "tenant_a");
assert_eq!(recent[1].target, "tenant_b");
}
#[tokio::test]
async fn load_recent_rehydrates_ring_from_disk() {
// Use a shared InMemory store so the second journal reads the first
// journal's writes — mirrors gateway restart semantics.
let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
{
let j1 = ErrorJournal {
log: Arc::new(AppendLog::new(store.clone(), JOURNAL_PREFIX).with_flush_threshold(1)),
ring: Arc::new(RwLock::new(VecDeque::with_capacity(RING_CAPACITY))),
};
j1.append(BucketErrorEvent::new_read("t", "k", "err")).await;
j1.flush().await.unwrap();
}
let j2 = ErrorJournal {
log: Arc::new(AppendLog::new(store.clone(), JOURNAL_PREFIX)),
ring: Arc::new(RwLock::new(VecDeque::with_capacity(RING_CAPACITY))),
};
assert_eq!(j2.recent(10).await.len(), 0, "fresh ring is empty");
let loaded = j2.load_recent().await.unwrap();
assert_eq!(loaded, 1);
assert_eq!(j2.recent(10).await[0].target, "t");
}
#[tokio::test]
async fn mark_rescued_updates_most_recent_match() {
let j = mk();
j.append(BucketErrorEvent::new_read("t1", "k", "err")).await;
j.append(BucketErrorEvent::new_read("t2", "k", "err")).await;
j.append(BucketErrorEvent::new_read("t1", "k", "err")).await;
j.mark_rescued_last("t1", "k").await;
let recent = j.recent(10).await;
assert!(!recent[0].rescued, "older t1 event not touched");
assert!(!recent[1].rescued, "t2 not touched");
assert!(recent[2].rescued, "most recent t1 was rescued");
}
#[tokio::test]
async fn health_flags_buckets_with_three_or_more_errors() {
let j = mk();
for _ in 0..3 { j.append(BucketErrorEvent::new_read("broken", "k", "e")).await; }
j.append(BucketErrorEvent::new_read("flaky", "k", "e")).await;
let h = j.health(60).await;
assert_eq!(h.total_errors, 4);
assert_eq!(h.per_bucket.get("broken"), Some(&3));
assert_eq!(h.per_bucket.get("flaky"), Some(&1));
assert!(h.unhealthy_buckets.contains(&"broken".to_string()));
assert!(!h.unhealthy_buckets.contains(&"flaky".to_string()));
}
#[tokio::test]
async fn filter_by_bucket_returns_chronological_order() {
let j = mk();
j.append(BucketErrorEvent::new_read("a", "k1", "e")).await;
j.append(BucketErrorEvent::new_read("b", "k2", "e")).await;
j.append(BucketErrorEvent::new_read("a", "k3", "e")).await;
let only_a = j.filter(Some("a"), None, 10).await;
assert_eq!(only_a.len(), 2);
assert_eq!(only_a[0].key, "k1", "oldest first");
assert_eq!(only_a[1].key, "k3");
}
}
+179
View File
@@ -0,0 +1,179 @@
/// HTTP surface for federation operator endpoints.
///
/// - `GET /buckets` — configured buckets with reachability status
/// - `GET /errors` — recent bucket op failures (filterable)
/// - `GET /bucket-health` — aggregated errors-per-bucket in the last 5 minutes
///
/// Mounted under `/storage` by the gateway. `/storage/health` is already
/// claimed by the existing storage router for service liveness, so the
/// aggregated bucket health lives at `/storage/bucket-health`.
use axum::{
Json, Router,
extract::{Path, Query, State},
http::StatusCode,
response::IntoResponse,
routing::{delete, get, post, put},
};
use chrono::{DateTime, Utc};
use serde::Deserialize;
use shared::config::BucketConfig;
use std::sync::Arc;
use crate::error_journal::BucketErrorEvent;
use crate::registry::{BucketInfo, BucketRegistry};
pub fn router(registry: Arc<BucketRegistry>) -> Router {
Router::new()
.route("/buckets", get(list_buckets).post(add_bucket))
.route("/buckets/{name}", delete(remove_bucket))
.route("/errors", get(list_errors))
.route("/errors/flush", post(flush_errors))
.route("/errors/compact", post(compact_errors))
.route("/bucket-health", get(get_health))
.route("/buckets/{bucket}/objects/{*key}", put(put_bucket_object))
.route("/buckets/{bucket}/objects/{*key}", get(get_bucket_object))
.with_state(registry)
}
/// Provision + register a bucket at runtime. Body is a `BucketConfig`.
///
/// Federation layer 2: profile buckets (`profile:alice`) and tenant
/// buckets can be added after startup without a service restart.
async fn add_bucket(
State(reg): State<Arc<BucketRegistry>>,
Json(bc): Json<BucketConfig>,
) -> impl IntoResponse {
// Safety net: local backends must land somewhere under the configured
// profile_root (for profile:*) or be otherwise explicitly rooted.
// Refuse empty / missing root to avoid "bucket created at ./" surprises.
if bc.backend == "local" {
match bc.root.as_deref() {
Some(r) if !r.trim().is_empty() => {}
_ => return Err((StatusCode::BAD_REQUEST,
"local bucket requires a non-empty 'root' path".to_string())),
}
}
match reg.add_bucket(bc).await {
Ok(info) => Ok((StatusCode::CREATED, Json(info))),
Err(e) => {
let code = if e.contains("already registered") {
StatusCode::CONFLICT
} else {
StatusCode::BAD_REQUEST
};
Err((code, e))
}
}
}
/// Unregister a bucket. Refused for primary / rescue / unknown names.
async fn remove_bucket(
State(reg): State<Arc<BucketRegistry>>,
Path(name): Path<String>,
) -> impl IntoResponse {
match reg.remove_bucket(&name) {
Ok(()) => Ok((StatusCode::OK, format!("unregistered: {name}"))),
Err(e) => {
let code = if e.contains("cannot remove") {
StatusCode::FORBIDDEN
} else if e.contains("not registered") {
StatusCode::NOT_FOUND
} else {
StatusCode::BAD_REQUEST
};
Err((code, e))
}
}
}
async fn list_buckets(State(reg): State<Arc<BucketRegistry>>) -> Json<Vec<BucketInfo>> {
Json(reg.list().await)
}
#[derive(Deserialize)]
struct ErrorQuery {
bucket: Option<String>,
since: Option<DateTime<Utc>>,
#[serde(default = "default_limit")]
limit: usize,
}
fn default_limit() -> usize { 50 }
async fn list_errors(
State(reg): State<Arc<BucketRegistry>>,
Query(q): Query<ErrorQuery>,
) -> Json<Vec<BucketErrorEvent>> {
Json(reg.journal().filter(q.bucket.as_deref(), q.since, q.limit).await)
}
#[derive(Deserialize)]
struct HealthQuery {
#[serde(default = "default_period")]
minutes: i64,
}
fn default_period() -> i64 { 5 }
async fn get_health(
State(reg): State<Arc<BucketRegistry>>,
Query(q): Query<HealthQuery>,
) -> impl IntoResponse {
Json(reg.journal().health(q.minutes).await)
}
/// Bucket-aware write. Hard-fails if the target bucket is unreachable;
/// never falls back to rescue for writes.
async fn put_bucket_object(
State(reg): State<Arc<BucketRegistry>>,
Path((bucket, key)): Path<(String, String)>,
body: bytes::Bytes,
) -> impl IntoResponse {
match reg.write_smart(&bucket, &key, body).await {
Ok(()) => (StatusCode::CREATED, format!("stored: {bucket}/{key}")).into_response(),
Err(e) => (StatusCode::SERVICE_UNAVAILABLE, e).into_response(),
}
}
/// Bucket-aware read. Falls through to the rescue bucket if the target
/// is unreachable or the key is missing; emits observability headers so
/// callers can detect the fallback.
async fn flush_errors(State(reg): State<Arc<BucketRegistry>>) -> impl IntoResponse {
match reg.journal().flush().await {
Ok(()) => (StatusCode::OK, "flushed").into_response(),
Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, e).into_response(),
}
}
async fn compact_errors(State(reg): State<Arc<BucketRegistry>>) -> impl IntoResponse {
match reg.journal().compact().await {
Ok(stats) => Json(stats).into_response(),
Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, e).into_response(),
}
}
async fn get_bucket_object(
State(reg): State<Arc<BucketRegistry>>,
Path((bucket, key)): Path<(String, String)>,
) -> impl IntoResponse {
match reg.read_smart(&bucket, &key).await {
Ok(outcome) => {
let mut resp = outcome.data.into_response();
if outcome.rescued {
let headers = resp.headers_mut();
headers.insert("x-lakehouse-rescue-used", "true".parse().unwrap());
headers.insert(
"x-lakehouse-original-bucket",
outcome.original_bucket.parse().unwrap(),
);
headers.insert(
"x-lakehouse-served-by",
outcome.served_by.parse().unwrap(),
);
}
resp
}
Err(e) => (StatusCode::NOT_FOUND, e).into_response(),
}
}
+7
View File
@@ -0,0 +1,7 @@
pub mod append_log;
pub mod backend;
pub mod error_journal;
pub mod federation_service;
pub mod ops;
pub mod registry;
pub mod service;
+30
View File
@@ -0,0 +1,30 @@
use bytes::Bytes;
use futures::TryStreamExt;
use object_store::{ObjectStore, path::Path};
use std::sync::Arc;
pub async fn put(store: &Arc<dyn ObjectStore>, key: &str, data: Bytes) -> Result<(), String> {
let path = Path::from(key);
store.put(&path, data.into()).await.map_err(|e| e.to_string())?;
Ok(())
}
pub async fn get(store: &Arc<dyn ObjectStore>, key: &str) -> Result<Bytes, String> {
let path = Path::from(key);
let result = store.get(&path).await.map_err(|e| e.to_string())?;
let bytes = result.bytes().await.map_err(|e| e.to_string())?;
Ok(bytes)
}
pub async fn delete(store: &Arc<dyn ObjectStore>, key: &str) -> Result<(), String> {
let path = Path::from(key);
store.delete(&path).await.map_err(|e| e.to_string())?;
Ok(())
}
pub async fn list(store: &Arc<dyn ObjectStore>, prefix: Option<&str>) -> Result<Vec<String>, String> {
let prefix_path = prefix.map(Path::from);
let stream = store.list(prefix_path.as_ref());
let items: Vec<_> = stream.try_collect().await.map_err(|e| e.to_string())?;
Ok(items.into_iter().map(|m| m.location.to_string()).collect())
}
+377
View File
@@ -0,0 +1,377 @@
/// Multi-backend bucket registry — the federation foundation.
///
/// Federation rule: every `ObjectRef` belongs to exactly one named bucket.
/// The registry resolves bucket names to `object_store` backends, handles
/// rescue-bucket fallback on read failure, writes every failure to the
/// error journal, and exposes a health summary for operators.
///
/// Existing call sites can keep using `ops::*` with `registry.get(name)`.
/// New bucket-aware call sites use `registry.read_smart` / `write_smart`
/// which handle fallback + journaling automatically.
use object_store::ObjectStore;
use object_store::local::LocalFileSystem;
use serde::Serialize;
use shared::config::{BucketConfig, StorageConfig};
use shared::secrets::{BucketCredentials, SecretsProvider};
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use crate::error_journal::{BucketErrorEvent, ErrorJournal};
/// A registered bucket — the store handle + its configuration.
pub struct BucketEntry {
pub name: String,
pub backend: String,
pub store: Arc<dyn ObjectStore>,
pub config: BucketConfig,
}
/// Read outcome — may have been rescued.
#[derive(Debug, Clone)]
pub struct ReadOutcome {
pub data: bytes::Bytes,
pub rescued: bool,
pub original_bucket: String,
pub served_by: String,
}
/// Summary entry for GET /storage/buckets.
#[derive(Debug, Clone, Serialize)]
pub struct BucketInfo {
pub name: String,
pub backend: String,
pub reachable: bool,
pub role: BucketRole,
}
#[derive(Debug, Clone, Serialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum BucketRole {
Primary,
Rescue,
Profile,
Tenant,
}
pub struct BucketRegistry {
/// RwLock because federation layer 2 adds runtime bucket lifecycle
/// (`POST /storage/buckets` / `DELETE /storage/buckets/{name}`).
/// Almost all accesses are reads — writes happen at provision /
/// deactivate time only.
buckets: RwLock<HashMap<String, Arc<BucketEntry>>>,
default: String,
rescue: Option<String>,
profile_root: String,
/// Held so runtime `add_bucket` can resolve secret_ref handles.
secrets: Arc<dyn SecretsProvider>,
journal: ErrorJournal,
}
impl BucketRegistry {
/// Build the registry from storage config + secrets provider.
/// Back-compat: if `buckets` is empty, synthesize a `primary` bucket from
/// the legacy `root` field so pre-federation configs keep working.
pub async fn from_config(
cfg: &StorageConfig,
secrets: Arc<dyn SecretsProvider>,
) -> Result<Self, String> {
let mut buckets: HashMap<String, Arc<BucketEntry>> = HashMap::new();
let bucket_configs: Vec<BucketConfig> = if cfg.buckets.is_empty() {
vec![BucketConfig {
name: "primary".to_string(),
backend: "local".to_string(),
root: Some(cfg.root.clone()),
bucket: None,
region: None,
endpoint: None,
secret_ref: None,
}]
} else {
cfg.buckets.clone()
};
for bc in bucket_configs {
let store = build_store(&bc, secrets.as_ref()).await?;
let entry = Arc::new(BucketEntry {
name: bc.name.clone(),
backend: bc.backend.clone(),
store,
config: bc.clone(),
});
buckets.insert(bc.name.clone(), entry);
}
// Ensure `primary` always exists — it's where error journals live.
if !buckets.contains_key("primary") {
return Err("no bucket named 'primary' configured — required as error-journal home".into());
}
// Rescue bucket is optional but, if named, must exist.
if let Some(r) = &cfg.rescue_bucket {
if !buckets.contains_key(r) {
return Err(format!("rescue_bucket '{r}' not found among configured buckets"));
}
}
let journal = ErrorJournal::new(buckets.get("primary").unwrap().store.clone());
let _ = journal.load_recent().await;
Ok(Self {
buckets: RwLock::new(buckets),
default: "primary".to_string(),
rescue: cfg.rescue_bucket.clone(),
profile_root: cfg.profile_root.clone(),
secrets,
journal,
})
}
pub fn default_name(&self) -> &str { &self.default }
pub fn rescue_name(&self) -> Option<&str> { self.rescue.as_deref() }
pub fn profile_root(&self) -> &str { &self.profile_root }
pub fn journal(&self) -> &ErrorJournal { &self.journal }
/// Resolve a bucket name to its object store. Existing call sites use
/// this as a drop-in replacement for the old single-store pattern.
///
/// Uses `std::sync::RwLock` — every caller holds the guard for
/// microseconds (clones an Arc and drops), so there's no async-
/// blocking concern. Never hold across an await.
pub fn get(&self, bucket: &str) -> Result<Arc<dyn ObjectStore>, String> {
let guard = self.buckets.read().expect("bucket registry lock poisoned");
guard
.get(bucket)
.map(|e| e.store.clone())
.ok_or_else(|| format!("unknown bucket: {bucket}"))
}
/// The default bucket's store — use for code paths that don't yet know
/// about buckets.
pub fn default_store(&self) -> Arc<dyn ObjectStore> {
let guard = self.buckets.read().expect("bucket registry lock poisoned");
guard.get(&self.default).unwrap().store.clone()
}
/// True if this bucket name is registered.
pub fn contains(&self, bucket: &str) -> bool {
let guard = self.buckets.read().expect("bucket registry lock poisoned");
guard.contains_key(bucket)
}
/// List all registered buckets. Checks reachability by doing a trivial
/// `list` with limit 1 on each.
pub async fn list(&self) -> Vec<BucketInfo> {
// Snapshot (name, backend, store_clone) under the lock, then probe
// outside — probing can be slow (network), don't hold the lock.
let snapshot: Vec<(String, String, Arc<dyn ObjectStore>)> = {
let guard = self.buckets.read().expect("bucket registry lock poisoned");
guard
.iter()
.map(|(n, e)| (n.clone(), e.backend.clone(), e.store.clone()))
.collect()
};
let mut out = Vec::with_capacity(snapshot.len());
for (name, backend, store) in snapshot {
let reachable = probe(&store).await;
let role = self.classify(&name);
out.push(BucketInfo { name, backend, reachable, role });
}
out.sort_by(|a, b| a.name.cmp(&b.name));
out
}
fn classify(&self, name: &str) -> BucketRole {
if name == self.default { BucketRole::Primary }
else if Some(name) == self.rescue.as_deref() { BucketRole::Rescue }
else if name.starts_with("profile:") { BucketRole::Profile }
else { BucketRole::Tenant }
}
/// Provision + register a bucket at runtime (federation layer 2).
/// Returns Err if the name is already registered. Local backends get
/// their root directory created automatically.
pub async fn add_bucket(&self, bc: BucketConfig) -> Result<BucketInfo, String> {
if self.contains(&bc.name) {
return Err(format!("bucket '{}' already registered", bc.name));
}
let store = build_store(&bc, self.secrets.as_ref()).await?;
let reachable = probe(&store).await;
let entry = Arc::new(BucketEntry {
name: bc.name.clone(),
backend: bc.backend.clone(),
store,
config: bc.clone(),
});
let role = self.classify(&bc.name);
let info = BucketInfo {
name: bc.name.clone(),
backend: bc.backend.clone(),
reachable,
role,
};
{
let mut guard = self.buckets.write().expect("bucket registry lock poisoned");
guard.insert(bc.name.clone(), entry);
}
tracing::info!("registered bucket '{}' backend={} reachable={}",
bc.name, bc.backend, reachable);
Ok(info)
}
/// Unregister a bucket. Refuses to remove `primary` or the configured
/// rescue bucket — those are load-bearing. Caller is responsible for
/// any final flush of bucket-local state before calling this.
pub fn remove_bucket(&self, name: &str) -> Result<(), String> {
if name == self.default {
return Err("cannot remove primary bucket".into());
}
if Some(name) == self.rescue.as_deref() {
return Err(format!("cannot remove rescue bucket '{name}'"));
}
let removed = {
let mut guard = self.buckets.write().expect("bucket registry lock poisoned");
guard.remove(name)
};
if removed.is_none() {
return Err(format!("bucket '{name}' not registered"));
}
tracing::info!("unregistered bucket '{}'", name);
Ok(())
}
/// Read with rescue-bucket fallback. If the target bucket fails and a
/// rescue is configured, retries against rescue. Records every failure
/// in the error journal.
pub async fn read_smart(&self, bucket: &str, key: &str) -> Result<ReadOutcome, String> {
let target_store = {
let guard = self.buckets.read().expect("bucket registry lock poisoned");
guard.get(bucket).map(|e| e.store.clone())
.ok_or_else(|| format!("unknown bucket: {bucket}"))?
};
match crate::ops::get(&target_store, key).await {
Ok(data) => Ok(ReadOutcome {
data, rescued: false,
original_bucket: bucket.to_string(),
served_by: bucket.to_string(),
}),
Err(err) => {
// Record failure regardless of what happens next.
self.journal.append(BucketErrorEvent::new_read(bucket, key, &err)).await;
// Try rescue, if any.
if let Some(rescue_name) = &self.rescue {
if rescue_name != bucket {
let rescue_store = {
let guard = self.buckets.read().expect("bucket registry lock poisoned");
guard.get(rescue_name).map(|e| e.store.clone())
};
if let Some(rescue_store) = rescue_store {
match crate::ops::get(&rescue_store, key).await {
Ok(data) => {
self.journal.mark_rescued_last(bucket, key).await;
return Ok(ReadOutcome {
data, rescued: true,
original_bucket: bucket.to_string(),
served_by: rescue_name.clone(),
});
}
Err(rescue_err) => {
return Err(format!(
"read '{key}' failed in '{bucket}' ({err}); rescue '{rescue_name}' also failed ({rescue_err})"
));
}
}
}
}
}
Err(format!("read '{key}' failed in '{bucket}': {err}"))
}
}
}
/// Write always goes to target. No rescue fallback for writes — writes
/// that silently vanish are the worst possible failure.
pub async fn write_smart(
&self,
bucket: &str,
key: &str,
data: bytes::Bytes,
) -> Result<(), String> {
let target_store = {
let guard = self.buckets.read().expect("bucket registry lock poisoned");
guard.get(bucket).map(|e| e.store.clone())
.ok_or_else(|| format!("unknown bucket: {bucket}"))?
};
match crate::ops::put(&target_store, key, data).await {
Ok(()) => Ok(()),
Err(err) => {
self.journal.append(BucketErrorEvent::new_write(bucket, key, &err)).await;
Err(format!("write '{key}' failed in '{bucket}': {err}"))
}
}
}
}
/// Trivial reachability check — try to list with limit 0.
async fn probe(store: &Arc<dyn ObjectStore>) -> bool {
use futures::StreamExt;
let mut stream = store.list(None);
// Pulling the first item confirms the store responds. Empty bucket = ok.
match stream.next().await {
Some(Ok(_)) => true,
None => true, // empty but reachable
Some(Err(_)) => false,
}
}
/// Build a concrete ObjectStore from a BucketConfig.
async fn build_store(
bc: &BucketConfig,
secrets: &dyn SecretsProvider,
) -> Result<Arc<dyn ObjectStore>, String> {
match bc.backend.as_str() {
"local" => {
let root = bc.root.as_deref()
.ok_or_else(|| format!("bucket '{}' is backend=local but has no root", bc.name))?;
std::fs::create_dir_all(root)
.map_err(|e| format!("create bucket dir '{root}': {e}"))?;
let fs = LocalFileSystem::new_with_prefix(root)
.map_err(|e| format!("init local bucket '{}': {e}", bc.name))?;
Ok(Arc::new(fs))
}
"s3" => {
let handle = bc.secret_ref.as_deref()
.ok_or_else(|| format!("s3 bucket '{}' has no secret_ref", bc.name))?;
let creds: BucketCredentials = secrets.resolve(handle).await?;
let s3_bucket = bc.bucket.as_deref()
.ok_or_else(|| format!("s3 bucket '{}' has no `bucket` name", bc.name))?;
let region = bc.region.as_deref().unwrap_or("us-east-1");
let mut builder = object_store::aws::AmazonS3Builder::new()
.with_bucket_name(s3_bucket)
.with_region(region)
.with_access_key_id(&creds.access_key)
.with_secret_access_key(&creds.secret_key);
if let Some(endpoint) = &bc.endpoint {
builder = builder.with_endpoint(endpoint);
// MinIO and other S3-compatible services often run on plain
// HTTP. object_store refuses HTTP by default — opt in when
// a custom endpoint is configured (TLS endpoints work either way).
if endpoint.starts_with("http://") {
builder = builder.with_allow_http(true);
}
}
let s3 = builder.build()
.map_err(|e| format!("init s3 bucket '{}': {e}", bc.name))?;
Ok(Arc::new(s3))
}
other => Err(format!("unknown backend '{other}' for bucket '{}'", bc.name)),
}
}
+75
View File
@@ -0,0 +1,75 @@
use axum::{
Router,
extract::{Path, Query, State},
http::StatusCode,
response::IntoResponse,
routing::{delete, get, put},
};
use bytes::Bytes;
use object_store::ObjectStore;
use serde::Deserialize;
use std::sync::Arc;
use crate::ops;
pub type StorageState = Arc<dyn ObjectStore>;
pub fn router(store: Arc<dyn ObjectStore>) -> Router {
Router::new()
.route("/health", get(health))
.route("/objects/{*key}", put(put_object))
.route("/objects/{*key}", get(get_object))
.route("/objects/{*key}", delete(delete_object))
.route("/objects", get(list_objects))
.with_state(store)
}
async fn health() -> &'static str {
"storaged ok"
}
async fn put_object(
State(store): State<StorageState>,
Path(key): Path<String>,
body: Bytes,
) -> impl IntoResponse {
match ops::put(&store, &key, body).await {
Ok(()) => (StatusCode::CREATED, format!("stored: {key}")),
Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, e),
}
}
async fn get_object(
State(store): State<StorageState>,
Path(key): Path<String>,
) -> impl IntoResponse {
match ops::get(&store, &key).await {
Ok(data) => Ok(data),
Err(e) => Err((StatusCode::NOT_FOUND, e)),
}
}
async fn delete_object(
State(store): State<StorageState>,
Path(key): Path<String>,
) -> impl IntoResponse {
match ops::delete(&store, &key).await {
Ok(()) => (StatusCode::OK, format!("deleted: {key}")),
Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, e),
}
}
#[derive(Deserialize)]
struct ListQuery {
prefix: Option<String>,
}
async fn list_objects(
State(store): State<StorageState>,
Query(q): Query<ListQuery>,
) -> impl IntoResponse {
match ops::list(&store, q.prefix.as_deref()).await {
Ok(keys) => Ok(axum::Json(keys)),
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
}
}
+11
View File
@@ -0,0 +1,11 @@
[package]
name = "truth"
version = "0.1.0"
edition = "2024"
[dependencies]
serde = { workspace = true }
serde_json = { workspace = true }
tokio = { workspace = true }
tracing = { workspace = true }
toml = { workspace = true }
+49
View File
@@ -0,0 +1,49 @@
//! DevOps task-class rules — scaffold for the long-horizon phase.
//!
//! Phase 42 PRD: "Terraform/Ansible rule shapes are scaffolded but
//! unpopulated until the long-horizon phase. Keeps the dispatcher
//! signature stable so no refactor needed later."
//!
//! This module is intentionally minimal. It registers no rules yet.
//! The `devops_rules` function exists so callers can compose it onto
//! a store (e.g. `devops_rules(staffing_rules(TruthStore::new()))`)
//! without branching on whether the DevOps phase has landed.
//!
//! When the long-horizon phase fleshes out the DevOps rule set, the
//! implementations drop in here — same `RuleCondition` primitives, same
//! `TruthStore::evaluate` contract, zero upstream refactor.
use crate::TruthStore;
/// Register DevOps rules on the store. Currently a no-op scaffold —
/// no rules are added. Safe to compose with other rule-set functions.
///
/// Planned task classes (not yet populated):
/// - `devops.terraform_plan` — `terraform validate` + pre-plan
/// sanity checks (no destroys without confirm flag, etc.)
/// - `devops.ansible_playbook` — `ansible-lint` + privileged-task
/// gates (no `become: true` on untagged hosts)
/// - `devops.shell_command` — whitelist / blocklist for
/// AI-generated shell invocations (covers what Phase 42
/// queryd SQL gate does for SQL — same idea, shell surface)
pub fn devops_rules(store: TruthStore) -> TruthStore {
// Intentionally empty. See module-level doc for the phased rollout.
store
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn devops_rules_is_a_noop_for_now() {
// Scaffold guarantee: composing devops_rules onto an empty
// store must not add any rules. Future long-horizon work will
// populate this and the assertion shifts to counting the
// expected additions.
let store = devops_rules(TruthStore::new());
assert_eq!(store.get_rules("devops.terraform_plan").len(), 0);
assert_eq!(store.get_rules("devops.ansible_playbook").len(), 0);
assert_eq!(store.get_rules("devops.shell_command").len(), 0);
}
}
+610
View File
@@ -0,0 +1,610 @@
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
pub mod staffing;
pub mod devops;
pub mod loader;
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct TruthRule {
pub id: String,
pub task_class: String,
pub description: String,
pub condition: RuleCondition,
pub action: RuleAction,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum RuleCondition {
Always,
FieldEquals { field: String, value: String },
FieldMismatch { field: String, value: String },
FieldEmpty { field: String },
FieldGreater { field: String, threshold: i64 },
// Case-insensitive substring scan — true if the field value contains
// ANY of `needles`. Added for SQL/command guards where rules of the
// form "sql must not contain DROP/DELETE/TRUNCATE" need to express
// enforcement as a passing precondition being absent.
FieldContainsAny { field: String, needles: Vec<String> },
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum RuleAction {
Pass,
Reject { message: String },
Redact { fields: Vec<String> },
Block { message: String },
}
#[derive(Default)]
pub struct TruthStore {
rules: HashMap<String, Vec<TruthRule>>,
}
impl TruthStore {
pub fn new() -> Self {
Self::default()
}
pub fn add_rule(&mut self, rule: TruthRule) {
self.rules
.entry(rule.task_class.clone())
.or_default()
.push(rule);
}
/// All rule IDs across every task class. Used by the file loader
/// to detect duplicate-ID collisions before registering new rules.
pub fn all_rule_ids(&self) -> std::collections::HashSet<String> {
self.rules
.values()
.flat_map(|v| v.iter().map(|r| r.id.clone()))
.collect()
}
pub fn get_rules(&self, task_class: &str) -> Vec<&TruthRule> {
self.rules
.get(task_class)
.map(|v| v.iter().collect())
.unwrap_or_default()
}
/// Legacy API: returns the list of actions registered for a task class
/// without evaluating conditions. Retained for backward compatibility
/// with callers that only want the action catalog. New callers should
/// prefer `evaluate()`, which actually walks `RuleCondition` against
/// a context and reports per-rule pass/fail.
pub fn check(&self, task_class: &str) -> Vec<RuleAction> {
let rules = self.get_rules(task_class);
rules
.into_iter()
.map(|r| r.action.clone())
.collect()
}
/// Evaluate every rule registered for `task_class` against `ctx`,
/// returning one `RuleOutcome` per rule. `passed = true` means the
/// rule's `condition` held; the rule's action is still attached so
/// callers can distinguish "passed and therefore no-op" (RuleAction::Pass)
/// from "passed and apply Redact". `passed = false` means the condition
/// failed — callers should treat the attached action as the enforcement
/// response (Reject/Block).
///
/// Fixed P42-001 (2026-04-23): previously `check()` returned all actions
/// unconditionally — the `RuleCondition` field was ignored. Now every
/// rule is actually walked against the provided context.
pub fn evaluate(&self, task_class: &str, ctx: &serde_json::Value) -> Vec<RuleOutcome> {
self.get_rules(task_class)
.into_iter()
.map(|r| RuleOutcome {
rule_id: r.id.clone(),
passed: evaluate_condition(&r.condition, ctx),
action: r.action.clone(),
})
.collect()
}
}
/// Result of evaluating one rule against a context. `passed` reports
/// whether the condition held; `action` is the rule's declared action
/// regardless (callers decide how to apply it based on `passed`).
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct RuleOutcome {
pub rule_id: String,
pub passed: bool,
pub action: RuleAction,
}
fn evaluate_condition(cond: &RuleCondition, ctx: &serde_json::Value) -> bool {
match cond {
RuleCondition::Always => true,
RuleCondition::FieldEquals { field, value } => {
field_as_string(ctx, field)
.map(|s| s == *value)
.unwrap_or(false)
}
RuleCondition::FieldMismatch { field, value } => {
field_as_string(ctx, field)
.map(|s| s != *value)
.unwrap_or(false)
}
RuleCondition::FieldEmpty { field } => {
match lookup(ctx, field) {
None => true,
Some(v) => v.is_null() || v.as_str().map(|s| s.is_empty()).unwrap_or(false),
}
}
RuleCondition::FieldGreater { field, threshold } => {
lookup(ctx, field)
.and_then(|v| v.as_i64().or_else(|| v.as_f64().map(|f| f as i64)))
.map(|n| n > *threshold)
.unwrap_or(false)
}
RuleCondition::FieldContainsAny { field, needles } => {
match field_as_string(ctx, field) {
None => false,
Some(s) => {
let haystack = s.to_ascii_lowercase();
needles.iter().any(|n| haystack.contains(&n.to_ascii_lowercase()))
}
}
}
}
}
/// Walk a dot-separated path through a serde_json::Value. `"worker.status"`
/// → `ctx["worker"]["status"]`. Returns None if any segment is missing or
/// a non-object is encountered mid-path.
fn lookup<'a>(ctx: &'a serde_json::Value, path: &str) -> Option<&'a serde_json::Value> {
let mut cur = ctx;
for seg in path.split('.') {
cur = cur.get(seg)?;
}
Some(cur)
}
fn field_as_string(ctx: &serde_json::Value, path: &str) -> Option<String> {
lookup(ctx, path).and_then(|v| match v {
serde_json::Value::String(s) => Some(s.clone()),
serde_json::Value::Bool(b) => Some(b.to_string()),
serde_json::Value::Number(n) => Some(n.to_string()),
_ => None,
})
}
/// Minimal SQL guard — rejects destructive verbs (DROP/TRUNCATE/DELETE).
/// queryd/src/service.rs loads this into its `QueryState` and evaluates
/// every `/sql` request against it before hitting the DataFusion engine.
/// This is the P42-002 enforcement point flagged across scrum iters 3-5
/// ("raw SQL forwarded without schema or policy gate").
///
/// Intentionally narrow: it's a safety net, not a full SQL parser. If
/// callers need richer AST-aware enforcement they should extend this with
/// structured rules rather than new needles.
pub fn sql_query_guard_store() -> TruthStore {
let mut store = TruthStore::new();
store.add_rule(TruthRule {
id: "no-destructive-sql".to_string(),
task_class: "sql_query".to_string(),
description: "SQL must not contain destructive verbs".to_string(),
condition: RuleCondition::FieldContainsAny {
field: "sql".to_string(),
needles: vec![
"drop table".to_string(),
"drop schema".to_string(),
"drop database".to_string(),
"truncate".to_string(),
"delete from".to_string(),
],
},
action: RuleAction::Reject {
message: "destructive SQL rejected by truth.sql_query_guard".to_string(),
},
});
store.add_rule(TruthRule {
id: "sql-not-empty".to_string(),
task_class: "sql_query".to_string(),
description: "SQL must not be empty".to_string(),
condition: RuleCondition::FieldEmpty {
field: "sql".to_string(),
},
action: RuleAction::Reject {
message: "empty SQL rejected".to_string(),
},
});
store
}
/// Phase 42 default store: staffing rules + DevOps scaffold composed
/// onto an empty TruthStore. Per the PRD: "Staffing rules ship first;
/// Terraform/Ansible rule shapes are scaffolded but unpopulated until
/// the long-horizon phase." The composition order is irrelevant here
/// (DevOps is empty) but preserved so the shape matches the PRD's
/// expected "compose on top" pattern.
///
/// Moved out of inline in-function rule registration (2026-04-24) to
/// land the Phase 42 module split the PRD called for: `staffing.rs` +
/// `devops.rs` each owns their task-class rule sets. Behavior unchanged
/// for existing callers.
pub fn default_truth_store() -> TruthStore {
devops::devops_rules(staffing::staffing_rules(TruthStore::new()))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn truth_store_new_is_empty() {
let store = TruthStore::new();
assert!(store.rules.is_empty());
}
#[test]
fn add_rule_inserts_into_correct_task_class() {
let mut store = TruthStore::new();
store.add_rule(TruthRule {
id: "test-rule".to_string(),
task_class: "test.task".to_string(),
description: "Test rule".to_string(),
condition: RuleCondition::Always,
action: RuleAction::Pass,
});
let rules = store.get_rules("test.task");
assert_eq!(rules.len(), 1);
assert_eq!(rules[0].id, "test-rule");
}
#[test]
fn get_rules_returns_empty_for_unknown_class() {
let store = TruthStore::new();
let rules = store.get_rules("unknown.class");
assert!(rules.is_empty());
}
#[test]
fn check_returns_actions_for_task_class() {
let mut store = TruthStore::new();
store.add_rule(TruthRule {
id: "a1".to_string(),
task_class: "test".to_string(),
description: "A1".to_string(),
condition: RuleCondition::Always,
action: RuleAction::Pass,
});
store.add_rule(TruthRule {
id: "a2".to_string(),
task_class: "test".to_string(),
description: "A2".to_string(),
condition: RuleCondition::Always,
action: RuleAction::Reject {
message: "test reject".to_string(),
},
});
let actions = store.check("test");
assert_eq!(actions.len(), 2);
}
#[test]
fn rule_condition_serialize_always() {
let cond = RuleCondition::Always;
let json = serde_json::to_string(&cond).unwrap();
assert!(json.contains(r#""type":"Always"#));
}
#[test]
fn rule_condition_serialize_field_equals() {
let cond = RuleCondition::FieldEquals {
field: "foo".to_string(),
value: "bar".to_string(),
};
let json = serde_json::to_string(&cond).unwrap();
assert!(json.contains(r#""type":"FieldEquals""#));
assert!(json.contains(r#""field":"foo""#));
assert!(json.contains(r#""value":"bar""#));
}
#[test]
fn rule_action_serialize_redact() {
let action = RuleAction::Redact {
fields: vec!["ssn".to_string()],
};
let json = serde_json::to_string(&action).unwrap();
assert!(json.contains(r#""type":"Redact""#));
assert!(json.contains("ssn"));
}
#[test]
fn rule_action_serialize_reject() {
let action = RuleAction::Reject {
message: "test".to_string(),
};
let json = serde_json::to_string(&action).unwrap();
assert!(json.contains(r#""type":"Reject""#));
}
#[test]
fn default_truth_store_has_staffing_rules() {
let store = default_truth_store();
let fill_rules = store.get_rules("staffing.fill");
assert!(!fill_rules.is_empty());
let any_rules = store.get_rules("staffing.any");
assert!(!any_rules.is_empty());
}
#[test]
fn multiple_rules_same_task_class() {
let mut store = TruthStore::new();
for i in 0..5 {
store.add_rule(TruthRule {
id: format!("rule-{}", i),
task_class: "test".to_string(),
description: format!("Rule {}", i),
condition: RuleCondition::Always,
action: RuleAction::Pass,
});
}
let rules = store.get_rules("test");
assert_eq!(rules.len(), 5);
}
#[test]
fn truth_rule_clone_preserves_data() {
let rule = TruthRule {
id: "clone-test".to_string(),
task_class: "clone.task".to_string(),
description: "Clone test".to_string(),
condition: RuleCondition::FieldEquals {
field: "x".to_string(),
value: "y".to_string(),
},
action: RuleAction::Block {
message: "blocked".to_string(),
},
};
let cloned = rule.clone();
assert_eq!(cloned.id, rule.id);
assert_eq!(cloned.condition, rule.condition);
assert_eq!(cloned.action, rule.action);
}
#[test]
fn field_greater_condition_parse() {
let json = r#"{"type":"FieldGreater","field":"count","threshold":10}"#;
let cond: RuleCondition = serde_json::from_str(json).unwrap();
match cond {
RuleCondition::FieldGreater { field, threshold } => {
assert_eq!(field, "count");
assert_eq!(threshold, 10);
}
_ => panic!("Expected FieldGreater"),
}
}
#[test]
fn block_action_blocks_with_message() {
let action = RuleAction::Block {
message: "Rate limited".to_string(),
};
let json = serde_json::to_string(&action).unwrap();
assert!(json.contains("Rate limited"));
}
#[test]
fn empty_store_check_returns_empty() {
let store = TruthStore::new();
let actions = store.check("empty.class");
assert!(actions.is_empty());
}
// ── P42-001 evaluate() tests — actually walk RuleCondition ──
fn fill_store() -> TruthStore {
let mut s = TruthStore::new();
s.add_rule(TruthRule {
id: "active".into(),
task_class: "t".into(),
description: "must be active".into(),
condition: RuleCondition::FieldEquals {
field: "worker.status".into(),
value: "active".into(),
},
action: RuleAction::Reject {
message: "worker not active".into(),
},
});
s.add_rule(TruthRule {
id: "deadline".into(),
task_class: "t".into(),
description: "deadline required".into(),
condition: RuleCondition::FieldEmpty {
field: "contract.deadline".into(),
},
action: RuleAction::Reject {
message: "missing deadline".into(),
},
});
s.add_rule(TruthRule {
id: "budget".into(),
task_class: "t".into(),
description: "budget positive".into(),
condition: RuleCondition::FieldGreater {
field: "contract.budget".into(),
threshold: 0,
},
action: RuleAction::Block {
message: "budget must be positive".into(),
},
});
s
}
#[test]
fn evaluate_field_equals_pass_on_match() {
let s = fill_store();
let ctx = serde_json::json!({"worker": {"status": "active"}});
let o = s.evaluate("t", &ctx);
let active = o.iter().find(|r| r.rule_id == "active").unwrap();
assert!(active.passed, "active condition should hold");
}
#[test]
fn evaluate_field_equals_fail_on_mismatch() {
let s = fill_store();
let ctx = serde_json::json!({"worker": {"status": "terminated"}});
let o = s.evaluate("t", &ctx);
let active = o.iter().find(|r| r.rule_id == "active").unwrap();
assert!(!active.passed, "terminated should fail active condition");
}
#[test]
fn evaluate_field_equals_fail_on_missing() {
let s = fill_store();
let ctx = serde_json::json!({});
let o = s.evaluate("t", &ctx);
let active = o.iter().find(|r| r.rule_id == "active").unwrap();
assert!(!active.passed, "missing worker.status should fail");
}
#[test]
fn evaluate_field_empty_pass_when_absent() {
let s = fill_store();
// FieldEmpty passes when the field is missing/null/empty string.
// Deadline rule says "field empty means action fires" — so passed=true
// here means the rule's condition held (deadline IS empty).
let ctx = serde_json::json!({});
let o = s.evaluate("t", &ctx);
let deadline = o.iter().find(|r| r.rule_id == "deadline").unwrap();
assert!(deadline.passed);
}
#[test]
fn evaluate_field_empty_fail_when_present() {
let s = fill_store();
let ctx = serde_json::json!({"contract": {"deadline": "2026-05-01"}});
let o = s.evaluate("t", &ctx);
let deadline = o.iter().find(|r| r.rule_id == "deadline").unwrap();
assert!(!deadline.passed, "non-empty deadline should fail FieldEmpty check");
}
#[test]
fn evaluate_field_greater_pass_and_fail() {
let s = fill_store();
let ctx_ok = serde_json::json!({"contract": {"budget": 100}});
let ctx_bad = serde_json::json!({"contract": {"budget": 0}});
let ok = s.evaluate("t", &ctx_ok);
let bad = s.evaluate("t", &ctx_bad);
assert!(ok.iter().find(|r| r.rule_id == "budget").unwrap().passed);
assert!(!bad.iter().find(|r| r.rule_id == "budget").unwrap().passed);
}
#[test]
fn evaluate_always_condition_passes_unconditionally() {
let mut s = TruthStore::new();
s.add_rule(TruthRule {
id: "always".into(),
task_class: "x".into(),
description: "".into(),
condition: RuleCondition::Always,
action: RuleAction::Pass,
});
let o = s.evaluate("x", &serde_json::json!(null));
assert!(o[0].passed);
}
#[test]
fn evaluate_preserves_action_regardless_of_outcome() {
let s = fill_store();
let ctx = serde_json::json!({"worker": {"status": "active"}});
let o = s.evaluate("t", &ctx);
let active = o.iter().find(|r| r.rule_id == "active").unwrap();
// Action is attached whether the rule passed or not — the consumer
// decides how to use it.
assert_eq!(
active.action,
RuleAction::Reject {
message: "worker not active".into()
}
);
}
#[test]
fn evaluate_on_unknown_task_class_returns_empty() {
let s = fill_store();
let o = s.evaluate("nonexistent", &serde_json::json!({}));
assert!(o.is_empty());
}
#[test]
fn check_still_returns_actions_unconditionally_for_back_compat() {
// Legacy API should still behave the same — no condition walking.
let s = fill_store();
let actions = s.check("t");
assert_eq!(actions.len(), 3, "check returns one action per rule regardless of condition");
}
fn sql_guard_store() -> TruthStore {
let mut s = TruthStore::new();
s.add_rule(TruthRule {
id: "no-destructive".into(),
task_class: "sql_query".into(),
description: "SQL must not contain destructive verbs".into(),
condition: RuleCondition::FieldContainsAny {
field: "sql".into(),
needles: vec![
"drop table".into(),
"drop schema".into(),
"truncate".into(),
"delete from".into(),
],
},
action: RuleAction::Reject {
message: "destructive SQL rejected".into(),
},
});
s
}
#[test]
fn field_contains_any_matches_case_insensitively() {
let s = sql_guard_store();
let ctx = serde_json::json!({"sql": "SELECT * FROM t; DROP TABLE users;"});
let o = s.evaluate("sql_query", &ctx);
assert!(o[0].passed, "condition holds when needle present (case-insensitive)");
}
#[test]
fn field_contains_any_is_false_when_no_needle_matches() {
let s = sql_guard_store();
let ctx = serde_json::json!({"sql": "SELECT count(*) FROM workers"});
let o = s.evaluate("sql_query", &ctx);
assert!(!o[0].passed, "benign SELECT should not match destructive needles");
}
#[test]
fn field_contains_any_false_when_field_missing() {
let s = sql_guard_store();
let ctx = serde_json::json!({});
let o = s.evaluate("sql_query", &ctx);
assert!(!o[0].passed, "missing field → condition cannot hold");
}
#[test]
fn field_contains_any_empty_needles_list_never_matches() {
let mut s = TruthStore::new();
s.add_rule(TruthRule {
id: "empty".into(),
task_class: "x".into(),
description: "".into(),
condition: RuleCondition::FieldContainsAny {
field: "sql".into(),
needles: vec![],
},
action: RuleAction::Pass,
});
let o = s.evaluate("x", &serde_json::json!({"sql": "anything"}));
assert!(!o[0].passed, "no needles → any::<bool> is false");
}
}
+187
View File
@@ -0,0 +1,187 @@
//! File-backed TruthRule loader (Phase 42 PRD).
//!
//! PRD: "truth/ dir at repo root — rule files, versioned in git."
//! This module walks a directory, parses every `*.toml` file it finds,
//! and registers the rules into a caller-supplied store. Rule IDs must
//! be unique across the combined set — duplicate-ID collisions are
//! load-time errors.
//!
//! The TOML format matches the shape at `truth/README.md`. The same
//! `RuleCondition` + `RuleAction` enums used by the in-code registrars
//! deserialize directly from `condition = { type = "FieldEquals", ... }`
//! thanks to `#[serde(tag = "type")]`.
use std::fs;
use std::path::Path;
use serde::Deserialize;
use crate::{TruthRule, TruthStore};
/// Deserialization wrapper — a TOML file is a list of [[rule]] blocks.
#[derive(Deserialize)]
struct RuleFile {
#[serde(default)]
rule: Vec<TruthRule>,
}
/// Load every `*.toml` file in `dir` and add its rules to `store`.
/// Returns the number of rules loaded across all files.
///
/// Errors:
/// - directory doesn't exist or can't be read
/// - any `.toml` file fails to parse
/// - any rule ID collides with an existing rule (same ID already
/// registered in the store)
///
/// Non-goals: recursive walk (flat dir only), hot reload (one-shot load).
pub fn load_from_dir(store: &mut TruthStore, dir: impl AsRef<Path>) -> Result<usize, String> {
let dir = dir.as_ref();
let entries = fs::read_dir(dir)
.map_err(|e| format!("read_dir {}: {e}", dir.display()))?;
let mut loaded_ids = store.all_rule_ids();
let mut count = 0usize;
let mut paths: Vec<_> = entries
.filter_map(|e| e.ok())
.map(|e| e.path())
.filter(|p| p.extension().and_then(|s| s.to_str()) == Some("toml"))
.collect();
// Deterministic order — alphabetical by filename. Matters when a
// cross-file ID collision happens; the earlier filename wins
// nothing (both error), but the error message is reproducible.
paths.sort();
for path in paths {
let raw = fs::read_to_string(&path)
.map_err(|e| format!("read {}: {e}", path.display()))?;
let file: RuleFile = toml::from_str(&raw)
.map_err(|e| format!("parse {}: {e}", path.display()))?;
for rule in file.rule {
if !loaded_ids.insert(rule.id.clone()) {
return Err(format!(
"duplicate rule id '{}' from {}",
rule.id,
path.display()
));
}
store.add_rule(rule);
count += 1;
}
}
Ok(count)
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Write;
fn write_file(dir: &Path, name: &str, content: &str) {
let path = dir.join(name);
let mut f = fs::File::create(&path).unwrap();
f.write_all(content.as_bytes()).unwrap();
}
#[test]
fn loads_rules_from_toml_files() {
let tmp = tempdir_for("loader_test");
write_file(&tmp, "a.toml", r#"
[[rule]]
id = "a-rule"
task_class = "test"
description = "test rule"
action = { type = "Pass" }
[rule.condition]
type = "Always"
"#);
let mut store = TruthStore::new();
let n = load_from_dir(&mut store, &tmp).unwrap();
assert_eq!(n, 1);
assert_eq!(store.get_rules("test").len(), 1);
let _ = fs::remove_dir_all(&tmp);
}
#[test]
fn rejects_duplicate_rule_ids() {
let tmp = tempdir_for("dup_ids");
write_file(&tmp, "a.toml", r#"
[[rule]]
id = "same"
task_class = "t"
description = ""
action = { type = "Pass" }
[rule.condition]
type = "Always"
"#);
write_file(&tmp, "b.toml", r#"
[[rule]]
id = "same"
task_class = "t"
description = ""
action = { type = "Pass" }
[rule.condition]
type = "Always"
"#);
let mut store = TruthStore::new();
let err = load_from_dir(&mut store, &tmp).unwrap_err();
assert!(err.contains("duplicate"), "got: {err}");
let _ = fs::remove_dir_all(&tmp);
}
#[test]
fn duplicate_with_in_code_rule_is_rejected() {
// Existing in-store IDs count as "already registered." Operator
// can't shadow an in-code rule by file without changing the ID.
let tmp = tempdir_for("dup_in_code");
write_file(&tmp, "conflict.toml", r#"
[[rule]]
id = "worker-active"
task_class = "staffing.fill"
description = "file attempt"
action = { type = "Pass" }
[rule.condition]
type = "Always"
"#);
// staffing_rules registers "worker-active"
let mut store = crate::staffing::staffing_rules(TruthStore::new());
let err = load_from_dir(&mut store, &tmp).unwrap_err();
assert!(err.contains("duplicate") && err.contains("worker-active"));
let _ = fs::remove_dir_all(&tmp);
}
#[test]
fn skips_non_toml_files() {
let tmp = tempdir_for("skip_non_toml");
write_file(&tmp, "a.toml", r#"
[[rule]]
id = "x"
task_class = "t"
description = ""
action = { type = "Pass" }
[rule.condition]
type = "Always"
"#);
write_file(&tmp, "README.md", "not a toml file");
let mut store = TruthStore::new();
let n = load_from_dir(&mut store, &tmp).unwrap();
assert_eq!(n, 1); // README.md ignored
let _ = fs::remove_dir_all(&tmp);
}
#[test]
fn missing_dir_returns_error() {
let mut store = TruthStore::new();
let err = load_from_dir(&mut store, "/nonexistent/path/here").unwrap_err();
assert!(err.contains("read_dir"));
}
fn tempdir_for(tag: &str) -> std::path::PathBuf {
let dir = std::env::temp_dir().join(format!("truth_loader_{}_{}", tag,
std::process::id()));
fs::create_dir_all(&dir).unwrap();
dir
}
}
+125
View File
@@ -0,0 +1,125 @@
//! Staffing task-class rules for the TruthStore.
//!
//! Phase 42 PRD: "Staffing rules ship first. Terraform/Ansible rule
//! shapes are scaffolded but unpopulated until the long-horizon phase."
//! This module owns the staffing rule set; `devops.rs` holds the
//! matching scaffold for the DevOps long-horizon.
//!
//! Rules registered here live under the task classes `staffing.fill`
//! (fill proposals), `staffing.rescue` (rescue escalations), and
//! `staffing.any` (rules that apply across all staffing task classes —
//! PII redaction being the canonical example).
//!
//! All rules are evaluated via the `TruthStore::evaluate` walk, which
//! pairs each rule's `RuleCondition` against a caller-supplied JSON
//! context and emits a `RuleOutcome { passed, action }` per rule.
//! Downstream enforcement (router gate, SQL gate, execution-loop gate)
//! decides how to apply the action — `Reject` / `Block` shortcircuit,
//! `Redact` mutates, `Pass` is informational.
use crate::{RuleAction, RuleCondition, TruthRule, TruthStore};
/// Register the staffing rule set on an existing store. Returns the
/// store for chaining if the caller wants to fold other rule sets on
/// top (e.g. `staffing_rules(devops_rules(TruthStore::new()))`).
pub fn staffing_rules(mut store: TruthStore) -> TruthStore {
store.add_rule(TruthRule {
id: "worker-active".to_string(),
task_class: "staffing.fill".to_string(),
description: "Worker must be active".to_string(),
condition: RuleCondition::FieldEquals {
field: "worker.status".to_string(),
value: "active".to_string(),
},
action: RuleAction::Pass,
});
store.add_rule(TruthRule {
id: "client-not-blacklisted".to_string(),
task_class: "staffing.fill".to_string(),
description: "Worker cannot be blacklisted for client".to_string(),
condition: RuleCondition::FieldEquals {
field: "worker.client_blacklisted".to_string(),
value: "false".to_string(),
},
action: RuleAction::Pass,
});
store.add_rule(TruthRule {
id: "deadline-required".to_string(),
task_class: "staffing.fill".to_string(),
description: "Contract must have deadline".to_string(),
condition: RuleCondition::FieldEmpty {
field: "contract.deadline".to_string(),
},
action: RuleAction::Reject {
message: "Contract deadline is required".to_string(),
},
});
store.add_rule(TruthRule {
id: "budget-required".to_string(),
task_class: "staffing.fill".to_string(),
description: "Budget must be non-negative".to_string(),
condition: RuleCondition::FieldGreater {
field: "contract.budget_per_hour_max".to_string(),
threshold: 0,
},
action: RuleAction::Pass,
});
store.add_rule(TruthRule {
id: "pii-redact".to_string(),
task_class: "staffing.any".to_string(),
description: "Redact PII before cloud calls".to_string(),
condition: RuleCondition::Always,
action: RuleAction::Redact {
fields: vec!["ssn".to_string(), "salary".to_string()],
},
});
store
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn staffing_rules_registers_five_rules() {
// 4 staffing.fill rules + 1 staffing.any rule = 5 total.
// Regression guard: if someone adds a rule to this module
// without updating the count, this test surfaces it.
let store = staffing_rules(TruthStore::new());
let fill = store.get_rules("staffing.fill").len();
let any = store.get_rules("staffing.any").len();
assert_eq!(fill, 4);
assert_eq!(any, 1);
}
#[test]
fn blacklisted_worker_fails_the_rule() {
let store = staffing_rules(TruthStore::new());
let ctx = serde_json::json!({
"worker": { "client_blacklisted": "true" }
});
let outcomes = store.evaluate("staffing.fill", &ctx);
let blk = outcomes.iter().find(|o| o.rule_id == "client-not-blacklisted").unwrap();
assert!(!blk.passed, "blacklisted worker must fail the rule");
}
#[test]
fn missing_deadline_fires_reject_via_empty_condition() {
let store = staffing_rules(TruthStore::new());
// FieldEmpty passes when the field is missing — and the rule's
// action is Reject, so enforcement should fire.
let ctx = serde_json::json!({});
let outcomes = store.evaluate("staffing.fill", &ctx);
let deadline = outcomes.iter().find(|o| o.rule_id == "deadline-required").unwrap();
assert!(deadline.passed);
match &deadline.action {
RuleAction::Reject { message } => assert!(message.contains("deadline")),
_ => panic!("expected Reject action"),
}
}
}
+12
View File
@@ -0,0 +1,12 @@
[package]
name = "ui"
version = "0.1.0"
edition = "2024"
[dependencies]
dioxus = { version = "0.7", features = ["web"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
tracing = "0.1"
reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] }
web-sys = { version = "0.3", features = ["Window", "Location"] }
+9
View File
@@ -0,0 +1,9 @@
[application]
name = "lakehouse"
asset_dir = "assets"
[web.app]
title = "Lakehouse"
[web.resource.dev]
style = ["/assets/style.css"]

Some files were not shown because too many files have changed in this diff Show More