Compare commits
10
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 rootandClaude Opus 4.6
|
4d64e4a618 | ||
|
rootandClaude Opus 4.6
|
3f19ad5d93 | ||
|
rootandClaude Opus 4.6
|
bd68b662b7 | ||
|
rootandClaude Opus 4.6
|
3cf000ccf3 | ||
|
rootandClaude Opus 4.6
|
1dcbe8e04c | ||
|
rootandClaude Opus 4.6
|
2189a6dbee | ||
|
rootandClaude Opus 4.6
|
bf033b2d9d | ||
|
rootandClaude Opus 4.6
|
0756075c6d | ||
|
rootandClaude Opus 4.6
|
8ca4ba5717 | ||
|
rootandClaude Opus 4.6
|
c49b30053b |
@@ -1,12 +0,0 @@
|
||||
{
|
||||
"mcpServers": {
|
||||
"lakehouse": {
|
||||
"command": "bun",
|
||||
"args": ["run", "/home/profit/lakehouse/mcp-server/index.ts"],
|
||||
"env": {
|
||||
"LAKEHOUSE_URL": "http://localhost:3100",
|
||||
"MCP_TRANSPORT": "stdio"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Generated
+403
-3767
File diff suppressed because it is too large
Load Diff
+1
-5
@@ -12,9 +12,6 @@ members = [
|
||||
"crates/journald",
|
||||
"crates/gateway",
|
||||
"crates/ui",
|
||||
"crates/lance-bench",
|
||||
"crates/vectord-lance",
|
||||
"crates/truth",
|
||||
]
|
||||
|
||||
[workspace.dependencies]
|
||||
@@ -48,5 +45,4 @@ csv = "1"
|
||||
lopdf = "0.35"
|
||||
encoding_rs = "0.8"
|
||||
instant-distance = "0.6"
|
||||
tokio-postgres = { version = "0.7", features = ["with-serde_json-1", "with-chrono-0_4", "with-uuid-1"] }
|
||||
mysql_async = { version = "0.34", default-features = false, features = ["minimal"] }
|
||||
tokio-postgres = { version = "0.7", features = ["with-serde_json-1", "with-chrono-0_4"] }
|
||||
|
||||
@@ -0,0 +1,183 @@
|
||||
# PRD: Observer Knowledge Base
|
||||
|
||||
**Status:** Active — extraction pipeline proven, needs multi-pass processing via LLM Team modes
|
||||
**Created:** 2026-03-28
|
||||
**Owner:** J
|
||||
**Branch:** `observer-kb`
|
||||
**Parent:** Lakehouse core (crates/)
|
||||
|
||||
---
|
||||
|
||||
## Problem
|
||||
|
||||
AI conversations generate massive amounts of knowledge that's lost to context windows, hallucination, and scope drift. Current approach (single-pass extraction into categories) produces disorganized, hard-to-validate output. Reading raw pages is not useful — the system needs to process, clean, structure, and curate knowledge like an editorial team.
|
||||
|
||||
## Vision
|
||||
|
||||
A publishing pipeline that takes raw AI conversation exports and produces well-curated, validated knowledge documents. The output should be something you hand to a team — organized, filtered, and clearly presented. Questions anchor the content, with branching connections showing what was asked, what was answered, and how reasoning was directed.
|
||||
|
||||
---
|
||||
|
||||
## Architecture
|
||||
|
||||
### Multi-Pass Processing via LLM Team Modes
|
||||
|
||||
The key insight: devop.live already has 20 analysis modes that process content through structured workflows. The Observer KB should route uploaded data through these modes automatically.
|
||||
|
||||
```
|
||||
Raw PDF Upload
|
||||
│
|
||||
▼ Mode: EXTRACT
|
||||
Structured facts, entities, relationships
|
||||
│
|
||||
▼ Mode: VALIDATOR
|
||||
Fact-check claims, score 1-10 confidence
|
||||
│
|
||||
▼ Mode: HALLUCINATION
|
||||
Verify each claim, flag likely hallucinations
|
||||
│
|
||||
▼ Mode: CONSENSUS
|
||||
Multiple passes iterate until extraction converges
|
||||
│
|
||||
▼ Mode: REDTEAM
|
||||
Find flaws in knowledge, patch them
|
||||
│
|
||||
▼ Mode: PIPELINE (custom)
|
||||
Clean → Q&A structure → topic grouping → value ranking
|
||||
│
|
||||
▼ RENDER
|
||||
Well-curated document: questions as anchors, branching connections,
|
||||
what was asked → answered → directed
|
||||
```
|
||||
|
||||
### Integration with devop.live (LLM Team)
|
||||
|
||||
The Observer KB calls the LLM Team API to run modes against extracted content:
|
||||
|
||||
| Mode | Purpose in KB |
|
||||
|---|---|
|
||||
| **extract** | Structured facts, entities, relationships from raw text |
|
||||
| **validator** | Fact-check extracted claims, confidence scores |
|
||||
| **hallucination** | Verify each factual claim independently |
|
||||
| **drift** | Process N times, find consistent vs uncertain content |
|
||||
| **redteam** | Find flaws in extracted knowledge, improve |
|
||||
| **consensus** | Iterate until multiple passes agree on correct extraction |
|
||||
| **pipeline** | Chain: clean → extract → validate → organize → render |
|
||||
| **research** | Generate follow-up questions, research in parallel |
|
||||
| **blindassembly** | Split content into parts, process independently, reassemble |
|
||||
|
||||
### Data Flow
|
||||
|
||||
```
|
||||
Upload Portal (Observer KB :3400)
|
||||
│
|
||||
├── PDF → Lakehouse ingestd (text extraction, Parquet storage)
|
||||
│
|
||||
▼
|
||||
Processing Pipeline
|
||||
│
|
||||
├── Call devop.live API (:5000) with each mode
|
||||
├── Each mode processes content through local Ollama models
|
||||
├── Results stored back in Lakehouse (Parquet + catalog)
|
||||
│
|
||||
▼
|
||||
Knowledge Store
|
||||
│
|
||||
├── knowledge_facts: verified claims with confidence scores
|
||||
├── knowledge_qa: question → answer pairs with context
|
||||
├── knowledge_topics: auto-generated categories
|
||||
├── knowledge_connections: how topics relate across conversations
|
||||
├── knowledge_contradictions: flagged conflicts between conversations
|
||||
├── knowledge_validated: claims that passed hallucination check
|
||||
│
|
||||
▼
|
||||
Rendering
|
||||
│
|
||||
├── Structured knowledge view (web UI)
|
||||
├── Exportable curated document (PDF/markdown)
|
||||
└── Question-anchored navigation (Q → A → follow-up → related)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Processing Pipeline Detail
|
||||
|
||||
### Pass 1: Clean + Structure
|
||||
- Strip PDF artifacts (headers, footers, page numbers)
|
||||
- Identify conversation turns (user/AI)
|
||||
- Fix encoding issues (ligatures, special chars)
|
||||
- Normalize formatting
|
||||
|
||||
### Pass 2: Q&A Extraction
|
||||
- Identify questions asked by user
|
||||
- Map each question to its answer
|
||||
- Track follow-up questions and how they branch
|
||||
- Build question→answer→follow-up chains
|
||||
|
||||
### Pass 3: Fact Extraction + Validation
|
||||
- Run EXTRACT mode: pull structured facts, entities, relationships
|
||||
- Run VALIDATOR mode: score each fact 1-10 confidence
|
||||
- Run HALLUCINATION mode: verify claims independently
|
||||
- Flag contradictions between conversations
|
||||
|
||||
### Pass 4: Organize + Curate
|
||||
- Run CONSENSUS mode: multiple passes converge on correct categorization
|
||||
- Auto-generate topic categories from content (not predefined)
|
||||
- Rank content by value (decisions > insights > discussion > noise)
|
||||
- Connect related content across conversations
|
||||
|
||||
### Pass 5: Render
|
||||
- Question-anchored view: each question is a node, answers branch from it
|
||||
- Topic-grouped view: knowledge organized by auto-generated categories
|
||||
- Timeline view: how thinking evolved across conversations
|
||||
- Export as clean markdown/PDF
|
||||
|
||||
---
|
||||
|
||||
## Current State (as of 2026-03-28)
|
||||
|
||||
### What's Built
|
||||
- PDF ingestion (8 PDFs, 477 pages extracted)
|
||||
- Single-pass knowledge extraction (496 topics, 300 decisions, 100 insights)
|
||||
- Semantic search via HNSW (27ms over 477 chunks)
|
||||
- RAG (question → search → LLM answer with sources)
|
||||
- 5-pass deep analysis (corpus overview, patterns, evolution, gaps, synthesis)
|
||||
- Observer KB web UI (knowledge-first, not page-first)
|
||||
- Upload portal with drag & drop
|
||||
|
||||
### What's Not Built Yet
|
||||
- Integration with devop.live LLM Team modes
|
||||
- Multi-pass processing pipeline (extract → validate → hallucination check)
|
||||
- Q&A chain extraction (question → answer → follow-up mapping)
|
||||
- Contradiction detection between conversations
|
||||
- Curated document export (PDF/markdown)
|
||||
- Question-anchored navigation
|
||||
- Auto-processing on upload (currently manual script)
|
||||
|
||||
### What Needs to Happen Next Session
|
||||
1. **Integrate with devop.live API** — call modes programmatically
|
||||
2. **Build the multi-pass pipeline** — extract → validate → hallucinate-check → organize
|
||||
3. **Q&A chain extraction** — questions as anchors, answers branching
|
||||
4. **Auto-process on upload** — no manual script, pipeline runs automatically
|
||||
5. **Render curated output** — readable, hand-it-to-someone quality
|
||||
|
||||
---
|
||||
|
||||
## Tech Stack
|
||||
|
||||
- **Storage:** Lakehouse storaged (Parquet on object storage)
|
||||
- **Catalog:** Lakehouse catalogd (metadata, PII, lineage)
|
||||
- **Query:** Lakehouse queryd (DataFusion SQL, hot cache, paged results)
|
||||
- **Vectors:** Lakehouse vectord (HNSW search, RAG)
|
||||
- **AI Processing:** devop.live LLM Team API (20 modes, multiple models)
|
||||
- **AI Runtime:** Ollama (local models — qwen2.5, mistral, gemma2, nomic-embed-text)
|
||||
- **UI:** Pure HTML/JS on :3400
|
||||
|
||||
---
|
||||
|
||||
## Non-Goals (for now)
|
||||
|
||||
- Real-time conversation ingestion (PDF exports only)
|
||||
- Multi-user collaboration
|
||||
- Cloud deployment
|
||||
- Editing extracted knowledge from UI (read-only for now)
|
||||
@@ -0,0 +1,286 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>Observer KB</title>
|
||||
<style>
|
||||
:root{--bg:#0a0a0f;--s1:#12121a;--s2:#1a1a28;--b:#2a2a3a;--t:#e0e0e8;--d:#888898;--a:#6c5ce7;--ok:#00d2a0;--err:#ff6b6b;--warn:#ffa94d;--f:'JetBrains Mono','Fira Code',monospace}
|
||||
*{margin:0;padding:0;box-sizing:border-box}body{font-family:var(--f);background:var(--bg);color:var(--t);font-size:13px;min-height:100vh}
|
||||
.app{max-width:1200px;margin:0 auto;padding:20px}
|
||||
header{display:flex;align-items:center;gap:20px;padding:16px 0;border-bottom:1px solid var(--b);margin-bottom:20px}
|
||||
header h1{font-size:20px;color:var(--a)}
|
||||
.tabs{display:flex;gap:2px}.tab{background:none;border:none;color:var(--d);font-family:var(--f);font-size:12px;padding:10px 16px;cursor:pointer;border-bottom:2px solid transparent}
|
||||
.tab:hover{color:var(--t)}.tab.active{color:var(--a);border-bottom-color:var(--a)}
|
||||
.hdr-stats{margin-left:auto;font-size:11px;color:var(--d)}
|
||||
.btn{background:var(--a);color:white;border:none;padding:8px 16px;border-radius:6px;font-family:var(--f);font-size:12px;cursor:pointer;font-weight:600}
|
||||
.btn:hover{opacity:.9}.btn:disabled{opacity:.5}.btn-sm{padding:5px 10px;font-size:11px}
|
||||
.hidden{display:none}
|
||||
.stats-bar{display:grid;grid-template-columns:repeat(auto-fill,minmax(130px,1fr));gap:8px;margin-bottom:20px}
|
||||
.stat{background:var(--s1);border:1px solid var(--b);border-radius:8px;padding:12px;text-align:center}
|
||||
.stat-val{font-size:22px;font-weight:700;color:var(--a)}.stat-label{font-size:10px;color:var(--d);text-transform:uppercase;margin-top:2px}
|
||||
.section{margin-bottom:24px}.section h2{font-size:13px;color:var(--d);text-transform:uppercase;letter-spacing:.08em;margin-bottom:10px;border-bottom:1px solid var(--b);padding-bottom:6px}
|
||||
.entry{background:var(--s1);border:1px solid var(--b);border-radius:8px;padding:14px;margin-bottom:8px}
|
||||
.badge{font-size:9px;text-transform:uppercase;letter-spacing:.08em;padding:2px 8px;border-radius:3px;display:inline-block;margin-bottom:6px;font-weight:600}
|
||||
.badge-fact{background:var(--a);color:white}.badge-entity{background:var(--ok);color:white}
|
||||
.badge-decision{background:#e67e22;color:white}.badge-insight{background:#3498db;color:white}
|
||||
.badge-question{background:var(--warn);color:#000}.badge-verified{background:var(--ok);color:white}
|
||||
.entry-text{font-size:12px;line-height:1.6}.entry-src{font-size:10px;color:var(--d);margin-top:4px}
|
||||
.cat-grid{display:grid;grid-template-columns:repeat(auto-fill,minmax(250px,1fr));gap:8px;margin-bottom:16px}
|
||||
.cat-card{background:var(--s1);border:1px solid var(--b);border-radius:8px;padding:14px;cursor:pointer;border-left:3px solid var(--a);transition:all .15s}
|
||||
.cat-card:hover{border-color:var(--a)}.cat-title{font-weight:600;font-size:13px}.cat-count{font-size:10px;color:var(--d);margin-top:2px}
|
||||
.search-row{display:flex;gap:8px;margin-bottom:16px}
|
||||
.search-row input{flex:1;background:var(--s1);border:2px solid var(--b);color:var(--t);font-family:var(--f);font-size:14px;padding:12px;border-radius:8px;outline:none}
|
||||
.search-row input:focus{border-color:var(--a)}
|
||||
.ask-input{width:100%;background:var(--s1);border:2px solid var(--b);color:var(--t);font-family:var(--f);font-size:13px;padding:12px;border-radius:8px;resize:vertical;min-height:80px;outline:none;margin-bottom:8px}
|
||||
.ask-input:focus{border-color:var(--a)}
|
||||
.ai-answer{background:var(--s2);border:1px solid var(--b);border-radius:8px;padding:16px;margin-top:12px;line-height:1.7;white-space:pre-wrap}
|
||||
.ai-meta{font-size:10px;color:var(--d);margin-top:6px}
|
||||
.upload-zone{border:2px dashed var(--b);border-radius:12px;padding:36px;text-align:center;cursor:pointer;margin-bottom:16px}
|
||||
.upload-zone:hover{border-color:var(--a);background:rgba(108,92,231,.03)}
|
||||
.upload-zone h3{margin-bottom:4px}.upload-zone p{color:var(--d);font-size:11px}
|
||||
.file-item{display:flex;align-items:center;justify-content:space-between;padding:8px 12px;border-bottom:1px solid var(--b);font-size:12px}
|
||||
.done{color:var(--ok)}.failed{color:var(--err)}.uploading{color:var(--a)}
|
||||
.result{background:var(--s2);border:1px solid var(--b);border-radius:8px;padding:12px;margin-bottom:6px;border-left:3px solid var(--a)}
|
||||
.result-score{font-size:10px;color:var(--a);margin-bottom:2px}.result-text{font-size:12px;line-height:1.5}
|
||||
.empty{color:var(--d);text-align:center;padding:32px}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="app">
|
||||
<header>
|
||||
<h1>OBSERVER KB</h1>
|
||||
<div class="tabs">
|
||||
<button class="tab active" onclick="showTab('knowledge')">Knowledge</button>
|
||||
<button class="tab" onclick="showTab('facts')">Facts & Entities</button>
|
||||
<button class="tab" onclick="showTab('explore')">Search</button>
|
||||
<button class="tab" onclick="showTab('ask')">Ask AI</button>
|
||||
<button class="tab" onclick="showTab('upload')">Ingest</button>
|
||||
</div>
|
||||
<div class="hdr-stats" id="hdr-stats"></div>
|
||||
</header>
|
||||
|
||||
<div id="tab-knowledge"></div>
|
||||
<div id="tab-facts" class="hidden"></div>
|
||||
<div id="tab-explore" class="hidden"></div>
|
||||
<div id="tab-ask" class="hidden"></div>
|
||||
<div id="tab-upload" class="hidden"></div>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
const API=`http://${location.hostname}:3100`;
|
||||
const $=id=>document.getElementById(id);
|
||||
async function q(sql){return(await(await fetch(API+'/query/sql',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({sql})})).json())}
|
||||
async function searchHNSW(query,k=10){try{return(await(await fetch(API+'/vectors/hnsw/search',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({index_name:'observer_kb',query,top_k:k})})).json()).results||[]}catch{return[]}}
|
||||
|
||||
function showTab(name){
|
||||
['knowledge','facts','explore','ask','upload'].forEach((t,i)=>{
|
||||
$('tab-'+t).classList.toggle('hidden',t!==name);
|
||||
document.querySelectorAll('.tab')[i].classList.toggle('active',t===name);
|
||||
});
|
||||
({knowledge:loadKnowledge,facts:loadFacts,explore:loadExplore,ask:loadAsk,upload:loadUpload})[name]();
|
||||
}
|
||||
|
||||
function makeEntry(type,text,source){
|
||||
const e=document.createElement('div');e.className='entry';
|
||||
const b=document.createElement('span');b.className='badge badge-'+type;b.textContent=type;
|
||||
const t=document.createElement('div');t.className='entry-text';t.textContent=text;
|
||||
const s=document.createElement('div');s.className='entry-src';s.textContent=source;
|
||||
e.append(b,t,s);return e;
|
||||
}
|
||||
|
||||
// === KNOWLEDGE (categories + decisions + insights) ===
|
||||
async function loadKnowledge(){
|
||||
const el=$('tab-knowledge');el.textContent='';
|
||||
try{
|
||||
const[topics,decisions,insights,questions,extractions]=await Promise.all([
|
||||
q("SELECT topic,COUNT(*) as n,COUNT(DISTINCT conversation) as c FROM knowledge_topics GROUP BY topic ORDER BY n DESC LIMIT 15"),
|
||||
q("SELECT conversation,page,decision FROM knowledge_decisions LIMIT 12"),
|
||||
q("SELECT conversation,page,insight FROM knowledge_insights LIMIT 8"),
|
||||
q("SELECT COUNT(*) as n FROM knowledge_questions"),
|
||||
q("SELECT type,COUNT(*) as n FROM processed_extractions GROUP BY type"),
|
||||
]);
|
||||
|
||||
const factCount=extractions.rows.find(r=>r.type==='fact')?.n||0;
|
||||
const entityCount=extractions.rows.find(r=>r.type==='entity')?.n||0;
|
||||
|
||||
// Stats
|
||||
const sb=document.createElement('div');sb.className='stats-bar';
|
||||
[{v:topics.row_count,l:'Categories'},{v:factCount,l:'Verified Facts'},{v:entityCount,l:'Entities'},{v:decisions.row_count,l:'Decisions'},{v:insights.row_count,l:'Insights'},{v:questions.rows[0]?.n||0,l:'Open Questions'}].forEach(s=>{
|
||||
const d=document.createElement('div');d.className='stat';
|
||||
const v=document.createElement('div');v.className='stat-val';v.textContent=s.v;
|
||||
const l=document.createElement('div');l.className='stat-label';l.textContent=s.l;
|
||||
d.append(v,l);sb.appendChild(d);
|
||||
});
|
||||
el.appendChild(sb);
|
||||
$('hdr-stats').textContent=factCount+' facts · '+entityCount+' entities · '+decisions.row_count+' decisions';
|
||||
|
||||
// Categories
|
||||
const cs=document.createElement('div');cs.className='section';
|
||||
cs.insertAdjacentHTML('beforeend','<h2>Auto-Generated Categories</h2>');
|
||||
const cg=document.createElement('div');cg.className='cat-grid';
|
||||
topics.rows.forEach(t=>{
|
||||
const c=document.createElement('div');c.className='cat-card';
|
||||
const n=document.createElement('div');n.className='cat-title';n.textContent=t.topic;
|
||||
const m=document.createElement('div');m.className='cat-count';m.textContent=t.n+' mentions · '+t.c+' conversations';
|
||||
c.append(n,m);c.onclick=()=>drillTopic(t.topic);cg.appendChild(c);
|
||||
});
|
||||
cs.appendChild(cg);el.appendChild(cs);
|
||||
|
||||
// Decisions
|
||||
const ds=document.createElement('div');ds.className='section';
|
||||
ds.insertAdjacentHTML('beforeend','<h2>Key Decisions</h2>');
|
||||
decisions.rows.forEach(d=>ds.appendChild(makeEntry('decision',d.decision,d.conversation+' p'+d.page)));
|
||||
el.appendChild(ds);
|
||||
|
||||
// Insights
|
||||
const is=document.createElement('div');is.className='section';
|
||||
is.insertAdjacentHTML('beforeend','<h2>Key Insights</h2>');
|
||||
insights.rows.forEach(d=>is.appendChild(makeEntry('insight',d.insight,d.conversation+' p'+d.page)));
|
||||
el.appendChild(is);
|
||||
}catch(e){el.textContent='Error: '+e}
|
||||
}
|
||||
|
||||
async function drillTopic(topic){
|
||||
const el=$('tab-knowledge');el.textContent='';
|
||||
const back=document.createElement('button');back.className='btn btn-sm';back.textContent='← Back';back.onclick=loadKnowledge;
|
||||
const h=document.createElement('h2');h.style.cssText='font-size:18px;margin:12px 0';h.textContent=topic;
|
||||
el.append(back,h);
|
||||
const results=await searchHNSW(topic,10);
|
||||
const s=document.createElement('div');s.className='section';
|
||||
s.insertAdjacentHTML('beforeend','<h2>Related Content</h2>');
|
||||
results.forEach(r=>{
|
||||
const e=document.createElement('div');e.className='result';
|
||||
const sc=document.createElement('div');sc.className='result-score';sc.textContent='['+r.score.toFixed(3)+'] '+r.doc_id;
|
||||
const t=document.createElement('div');t.className='result-text';t.textContent=r.chunk_text.substring(0,300);
|
||||
e.append(sc,t);s.appendChild(e);
|
||||
});
|
||||
el.appendChild(s);
|
||||
}
|
||||
|
||||
// === FACTS & ENTITIES (from LLM Team extract mode) ===
|
||||
async function loadFacts(){
|
||||
const el=$('tab-facts');el.textContent='';
|
||||
try{
|
||||
const[facts,entities,verified]=await Promise.all([
|
||||
q("SELECT conversation,page,content FROM processed_extractions WHERE type='fact' ORDER BY conversation,page LIMIT 30"),
|
||||
q("SELECT conversation,content,COUNT(*) as n FROM processed_extractions WHERE type='entity' GROUP BY conversation,content ORDER BY n DESC LIMIT 30"),
|
||||
q("SELECT conversation,page,content FROM processed_validated LIMIT 10"),
|
||||
]);
|
||||
|
||||
const fs=document.createElement('div');fs.className='section';
|
||||
fs.insertAdjacentHTML('beforeend','<h2>Extracted Facts ('+facts.row_count+' total — via LLM Team Extract Mode)</h2>');
|
||||
facts.rows.forEach(r=>fs.appendChild(makeEntry('fact',r.content,r.conversation+' p'+r.page)));
|
||||
el.appendChild(fs);
|
||||
|
||||
const es=document.createElement('div');es.className='section';
|
||||
es.insertAdjacentHTML('beforeend','<h2>Entities ('+entities.row_count+' unique)</h2>');
|
||||
const eg=document.createElement('div');eg.style.cssText='display:flex;flex-wrap:wrap;gap:6px';
|
||||
entities.rows.forEach(r=>{
|
||||
const chip=document.createElement('span');
|
||||
chip.style.cssText='background:var(--s2);border:1px solid var(--b);padding:4px 10px;border-radius:4px;font-size:11px;cursor:pointer';
|
||||
chip.textContent=r.content+(r.n>1?' ('+r.n+')':'');
|
||||
chip.onclick=()=>{showTab('explore');setTimeout(()=>{document.getElementById('searchInput').value=r.content;doSearch()},100)};
|
||||
eg.appendChild(chip);
|
||||
});
|
||||
es.appendChild(eg);el.appendChild(es);
|
||||
|
||||
if(verified.rows.length){
|
||||
const vs=document.createElement('div');vs.className='section';
|
||||
vs.insertAdjacentHTML('beforeend','<h2>Verified by Secondary Model</h2>');
|
||||
verified.rows.forEach(r=>vs.appendChild(makeEntry('verified',r.content,r.conversation+' p'+r.page)));
|
||||
el.appendChild(vs);
|
||||
}
|
||||
}catch(e){el.textContent='Error: '+e}
|
||||
}
|
||||
|
||||
// === SEARCH ===
|
||||
function loadExplore(){
|
||||
const el=$('tab-explore');el.textContent='';
|
||||
const row=document.createElement('div');row.className='search-row';
|
||||
const input=document.createElement('input');input.id='searchInput';input.placeholder='Search by meaning across all knowledge...';
|
||||
const btn=document.createElement('button');btn.className='btn';btn.textContent='Search';
|
||||
btn.onclick=doSearch;input.onkeydown=e=>{if(e.key==='Enter')doSearch()};
|
||||
row.append(input,btn);
|
||||
const res=document.createElement('div');res.id='searchResults';
|
||||
const hint=document.createElement('div');hint.className='empty';hint.textContent='Try: "object storage architecture", "data ingestion trade-offs", "CRM options"';
|
||||
el.append(row,res,hint);
|
||||
}
|
||||
async function doSearch(){
|
||||
const query=document.getElementById('searchInput').value;if(!query)return;
|
||||
const el=$('searchResults');el.textContent='searching...';
|
||||
const results=await searchHNSW(query,15);
|
||||
el.textContent='';
|
||||
if(!results.length){el.textContent='No results';return}
|
||||
results.forEach(r=>{
|
||||
const e=document.createElement('div');e.className='result';
|
||||
const s=document.createElement('div');s.className='result-score';s.textContent='['+r.score.toFixed(3)+'] '+r.doc_id.replace(/_/g,' ');
|
||||
const t=document.createElement('div');t.className='result-text';t.textContent=r.chunk_text.substring(0,400);
|
||||
e.append(s,t);el.appendChild(e);
|
||||
});
|
||||
}
|
||||
|
||||
// === ASK AI ===
|
||||
function loadAsk(){
|
||||
const el=$('tab-ask');el.textContent='';
|
||||
const ta=document.createElement('textarea');ta.className='ask-input';ta.id='askInput';
|
||||
ta.placeholder='Ask a question about your knowledge base...\n\nExamples:\n- What architecture decisions were made about Iceberg?\n- Compare Spark vs Flink approaches\n- What approaches were abandoned and why?';
|
||||
const btn=document.createElement('button');btn.className='btn';btn.id='askBtn';btn.textContent='Ask AI';
|
||||
btn.onclick=doAsk;
|
||||
const resp=document.createElement('div');resp.id='askResp';
|
||||
el.append(ta,btn,resp);
|
||||
}
|
||||
async function doAsk(){
|
||||
const question=$('askInput').value;if(!question)return;
|
||||
$('askBtn').disabled=true;$('askBtn').textContent='thinking...';
|
||||
const el=$('askResp');el.textContent='searching knowledge...';
|
||||
try{
|
||||
const r=await(await fetch(API+'/vectors/rag',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({index_name:'observer_kb',question,top_k:5})})).json();
|
||||
el.textContent='';
|
||||
const ans=document.createElement('div');ans.className='ai-answer';ans.textContent=r.answer;
|
||||
const meta=document.createElement('div');meta.className='ai-meta';meta.textContent='Model: '+r.model+(r.tokens_generated?' · '+r.tokens_generated+' tokens':'');
|
||||
el.append(ans,meta);
|
||||
if(r.sources?.length){
|
||||
const sh=document.createElement('h2');sh.style.cssText='font-size:11px;color:var(--d);margin-top:14px';sh.textContent='SOURCES';el.appendChild(sh);
|
||||
r.sources.forEach(s=>{
|
||||
const e=document.createElement('div');e.className='result';e.style.marginTop='4px';
|
||||
const sc=document.createElement('div');sc.className='result-score';sc.textContent='['+s.score.toFixed(3)+'] '+s.doc_id;
|
||||
const t=document.createElement('div');t.className='result-text';t.style.fontSize='11px';t.textContent=s.chunk_text.substring(0,200);
|
||||
e.append(sc,t);el.appendChild(e);
|
||||
});
|
||||
}
|
||||
}catch(e){el.textContent='Error: '+e}
|
||||
$('askBtn').disabled=false;$('askBtn').textContent='Ask AI';
|
||||
}
|
||||
|
||||
// === UPLOAD ===
|
||||
function loadUpload(){
|
||||
const el=$('tab-upload');el.textContent='';
|
||||
const zone=document.createElement('div');zone.className='upload-zone';
|
||||
zone.onclick=()=>{const i=document.createElement('input');i.type='file';i.multiple=true;i.accept='.pdf,.txt,.csv,.json';i.onchange=e=>[...e.target.files].forEach(uploadFile);i.click()};
|
||||
zone.ondragover=e=>{e.preventDefault();zone.style.borderColor='var(--a)'};
|
||||
zone.ondragleave=()=>zone.style.borderColor='';
|
||||
zone.ondrop=e=>{e.preventDefault();zone.style.borderColor='';[...e.dataTransfer.files].forEach(uploadFile)};
|
||||
const h=document.createElement('h3');h.textContent='Drop files to ingest';
|
||||
const p=document.createElement('p');p.textContent='PDF, CSV, JSON, Text — auto-extracted and indexed';
|
||||
zone.append(h,p);
|
||||
const list=document.createElement('div');list.id='upList';
|
||||
el.append(zone,list);
|
||||
}
|
||||
function uploadFile(file){
|
||||
const item=document.createElement('div');item.className='file-item';
|
||||
const name=document.createElement('span');name.textContent=file.name;
|
||||
const status=document.createElement('span');status.className='uploading';status.textContent='uploading...';
|
||||
item.append(name,status);$('upList').prepend(item);
|
||||
const fd=new FormData();fd.append('file',file);
|
||||
const xhr=new XMLHttpRequest();xhr.open('POST',API+'/ingest/file');
|
||||
xhr.upload.onprogress=e=>{if(e.lengthComputable)status.textContent=Math.round(e.loaded/e.total*100)+'%'};
|
||||
xhr.onload=()=>{status.className=xhr.status<300?'done':'failed';status.textContent=xhr.status<300?'ingested':'failed'};
|
||||
xhr.onerror=()=>{status.className='failed';status.textContent='error'};
|
||||
xhr.send(fd);
|
||||
}
|
||||
|
||||
loadKnowledge();
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
@@ -1,60 +0,0 @@
|
||||
# Lakehouse Claim Auditor
|
||||
|
||||
A Bun sub-agent that watches open PRs on Gitea, reads the ship-claims
|
||||
in commit messages and PR bodies, and **hard-blocks** merges when the
|
||||
code doesn't back the claim.
|
||||
|
||||
Rationale: when "compiles + one curl works" gets called "phase shipped,"
|
||||
placeholder code accumulates. This auditor runs every 90s, fetches
|
||||
each open PR, and subjects it to four checks:
|
||||
|
||||
1. **Static diff** — grep/parse looking for placeholder patterns
|
||||
2. **Dynamic** — runs the never-before-executed hybrid test fixture
|
||||
3. **Cloud inference** — asks `gpt-oss:120b` via `/v1/chat` to
|
||||
identify gaps in the diff
|
||||
4. **KB query** — looks up `data/_kb/` + observer for prior failure
|
||||
patterns on similar claims
|
||||
|
||||
Verdict is assembled, posted to Gitea as:
|
||||
- A **failing commit status** (hard block — branch protection
|
||||
prevents merge)
|
||||
- A **review comment** explaining every finding
|
||||
|
||||
## Run manually
|
||||
|
||||
```bash
|
||||
cd /home/profit/lakehouse
|
||||
bun run auditor/index.ts
|
||||
```
|
||||
|
||||
Defaults: polls every 90s, stops on `auditor.paused` file present.
|
||||
|
||||
## State
|
||||
|
||||
- `data/_auditor/state.json` — last-audited head SHA per PR
|
||||
- `data/_auditor/verdicts/{pr}-{sha}.json` — per-run verdict record
|
||||
- `data/_kb/audit_lessons.jsonl` — one row per block/warn finding,
|
||||
path-agnostic signature for dedup. Tailed by kb_query on each audit
|
||||
to surface recurring patterns (2+ distinct PRs with same signature
|
||||
→ info, 3-4 → warn, 5+ → block). This is how the auditor learns.
|
||||
- `data/_kb/scrum_reviews.jsonl` — scrum-master per-file reviews. If
|
||||
a file in the current PR has been scrum-reviewed, kb_query surfaces
|
||||
the review as a finding with the accepted model and attempt count.
|
||||
|
||||
## Where YOU edit
|
||||
|
||||
`auditor/policy.ts` — the verdict assembler. Controls which findings
|
||||
block vs warn vs inform. All other code is mechanical: fetching,
|
||||
running checks, posting to Gitea.
|
||||
|
||||
## Hard-block mechanism
|
||||
|
||||
1. Commit status is posted as `failure` with context `lakehouse/auditor`
|
||||
2. If `main` branch protection requires `lakehouse/auditor` status
|
||||
to pass, Gitea prevents merge
|
||||
3. When code is fixed and re-audit passes, status flips to `success`,
|
||||
merge unblocks
|
||||
|
||||
Enable branch protection (one-time, via Gitea UI or API):
|
||||
- `POST /repos/profit/lakehouse/branch_protections`
|
||||
- `{"branch_name": "main", "required_status_checks": {"contexts": ["lakehouse/auditor"]}}`
|
||||
@@ -1,222 +0,0 @@
|
||||
// Orchestrator — runs all four checks on a PR, assembles a verdict,
|
||||
// posts to Gitea. This is task #8's integration layer; the poller
|
||||
// (task #9) calls this once per PR on every fresh head SHA.
|
||||
//
|
||||
// Hard-block mechanism: commit status posted with state="failure"
|
||||
// and context="lakehouse/auditor". If `main` branch protection
|
||||
// requires that context to pass, merge is physically impossible
|
||||
// until the auditor re-audits a fixed commit and flips the status
|
||||
// to "success".
|
||||
//
|
||||
// Human-readable reasoning: posted as a PR issue comment (not a
|
||||
// review — reviews have self-review restrictions on Gitea and the
|
||||
// auditor currently uses the same PAT as the PR author).
|
||||
|
||||
import { readFile, writeFile, mkdir, appendFile } from "node:fs/promises";
|
||||
import { createHash } from "node:crypto";
|
||||
import { join } from "node:path";
|
||||
import type { PrSnapshot, Verdict, Finding } from "./types.ts";
|
||||
import { getPrDiff, postCommitStatus, postIssueComment } from "./gitea.ts";
|
||||
import { parseClaims } from "./claim_parser.ts";
|
||||
import { assembleVerdict } from "./policy.ts";
|
||||
import { runStaticCheck } from "./checks/static.ts";
|
||||
import { runDynamicCheck } from "./checks/dynamic.ts";
|
||||
import { runInferenceCheck } from "./checks/inference.ts";
|
||||
import { runKbCheck } from "./checks/kb_query.ts";
|
||||
|
||||
const VERDICTS_DIR = "/home/profit/lakehouse/data/_auditor/verdicts";
|
||||
// Playbook for audit findings — one row per block/warn finding from a
|
||||
// verdict. kb_query tails this next audit and escalates recurrences.
|
||||
// Structured as JSONL so it's cheap to append and cheap to tail.
|
||||
const AUDIT_LESSONS_JSONL = "/home/profit/lakehouse/data/_kb/audit_lessons.jsonl";
|
||||
|
||||
export interface AuditOptions {
|
||||
// Skip the cloud inference call (fast path for iteration). Default false.
|
||||
skip_inference?: boolean;
|
||||
// Skip the dynamic check (avoid running the hybrid fixture every PR,
|
||||
// since it hits live services and mutates playbook state). Default false
|
||||
// on `main`-branch-target PRs, true when auditing feature branches
|
||||
// where the fixture would pollute state. Caller decides.
|
||||
skip_dynamic?: boolean;
|
||||
// Skip Gitea posting — useful for dry-runs / local testing.
|
||||
// Default false.
|
||||
dry_run?: boolean;
|
||||
}
|
||||
|
||||
export async function auditPr(pr: PrSnapshot, opts: AuditOptions = {}): Promise<Verdict> {
|
||||
const t0 = Date.now();
|
||||
const diff = await getPrDiff(pr.number);
|
||||
const { claims } = parseClaims(pr);
|
||||
|
||||
// Run checks in parallel where they don't share mutable state.
|
||||
// Static + kb_query + inference are all read-only. Dynamic mutates
|
||||
// playbook state (nonce-scoped per run, but still live) so if
|
||||
// skip_dynamic is false we still run it in parallel — the mutation
|
||||
// is namespaced.
|
||||
const [staticFindings, dynamicFindings, inferenceFindings, kbFindings] = await Promise.all([
|
||||
runStaticCheck(diff),
|
||||
opts.skip_dynamic ? Promise.resolve(stubFinding("dynamic", "skipped by options")) : runDynamicCheck(),
|
||||
opts.skip_inference ? Promise.resolve(stubFinding("inference", "skipped by options")) : runInferenceCheck(claims, diff, { pr_number: pr.number, head_sha: pr.head_sha }),
|
||||
runKbCheck(claims, pr.files.map(f => f.path)),
|
||||
]);
|
||||
|
||||
const allFindings: Finding[] = [
|
||||
...staticFindings,
|
||||
...dynamicFindings,
|
||||
...inferenceFindings,
|
||||
...kbFindings,
|
||||
];
|
||||
|
||||
const duration_ms = Date.now() - t0;
|
||||
const metrics = {
|
||||
audit_duration_ms: duration_ms,
|
||||
findings_total: allFindings.length,
|
||||
findings_block: allFindings.filter(f => f.severity === "block").length,
|
||||
findings_warn: allFindings.filter(f => f.severity === "warn").length,
|
||||
findings_info: allFindings.filter(f => f.severity === "info").length,
|
||||
claims_strong: claims.filter(c => c.strength === "strong").length,
|
||||
claims_moderate: claims.filter(c => c.strength === "moderate").length,
|
||||
claims_weak: claims.filter(c => c.strength === "weak").length,
|
||||
claims_empirical: claims.filter(c => c.strength === "empirical").length,
|
||||
claims_total: claims.length,
|
||||
diff_bytes: diff.length,
|
||||
};
|
||||
|
||||
const verdict = assembleVerdict(allFindings, metrics, pr.number, pr.head_sha);
|
||||
|
||||
await persistVerdict(verdict);
|
||||
|
||||
// Feedback loop — every block/warn finding becomes a row in
|
||||
// audit_lessons.jsonl, dedup-keyed by (check, normalized-summary).
|
||||
// The next audit's kb_query reads these and escalates recurring
|
||||
// findings so we don't lose the "this pattern has been flagged
|
||||
// before" signal across runs. Fire-and-forget; failure here must
|
||||
// not break the audit.
|
||||
appendAuditLessons(verdict).catch(e =>
|
||||
console.error(`[audit] audit_lessons append failed: ${(e as Error).message}`));
|
||||
|
||||
if (!opts.dry_run) {
|
||||
await postToGitea(verdict);
|
||||
}
|
||||
|
||||
return verdict;
|
||||
}
|
||||
|
||||
// Normalizes a finding summary for dedup: strips path-specific tails
|
||||
// ("in path/to/file.ts" → "in <file>"), line numbers, and long
|
||||
// commit-hash snippets. The goal is: the SAME class of finding on
|
||||
// DIFFERENT files should share a signature, so we can measure
|
||||
// "this pattern keeps showing up."
|
||||
function normalizedSignature(f: Finding): string {
|
||||
const summary = String(f.summary)
|
||||
.replace(/\bin\s+\S+\.(ts|rs|js|py|md)\b/gi, "in <file>")
|
||||
.replace(/:\+?\d+\b/g, ":<line>")
|
||||
.replace(/[0-9a-f]{8,}/gi, "<hash>")
|
||||
.replace(/\s+/g, " ")
|
||||
.trim()
|
||||
.slice(0, 240);
|
||||
const src = `${f.check}::${f.severity}::${summary}`;
|
||||
return createHash("sha256").update(src).digest("hex").slice(0, 16);
|
||||
}
|
||||
|
||||
async function appendAuditLessons(v: Verdict): Promise<void> {
|
||||
const actionable = v.findings.filter(f => f.severity === "block" || f.severity === "warn");
|
||||
if (actionable.length === 0) return;
|
||||
await mkdir(join(AUDIT_LESSONS_JSONL, ".."), { recursive: true });
|
||||
const rows: string[] = [];
|
||||
for (const f of actionable) {
|
||||
rows.push(JSON.stringify({
|
||||
signature: normalizedSignature(f),
|
||||
check: f.check,
|
||||
severity: f.severity,
|
||||
summary: f.summary,
|
||||
pr_number: v.pr_number,
|
||||
head_sha: v.head_sha,
|
||||
audited_at: v.audited_at,
|
||||
}));
|
||||
}
|
||||
await appendFile(AUDIT_LESSONS_JSONL, rows.join("\n") + "\n");
|
||||
}
|
||||
|
||||
async function persistVerdict(v: Verdict): Promise<void> {
|
||||
await mkdir(VERDICTS_DIR, { recursive: true });
|
||||
const filename = `${v.pr_number}-${v.head_sha.slice(0, 12)}.json`;
|
||||
await writeFile(join(VERDICTS_DIR, filename), JSON.stringify(v, null, 2));
|
||||
}
|
||||
|
||||
export async function postToGitea(v: Verdict): Promise<void> {
|
||||
// 1. Commit status — the hard block signal (if branch protection
|
||||
// is configured to require lakehouse/auditor on main).
|
||||
const state = v.overall === "approve" ? "success" : "failure";
|
||||
await postCommitStatus({
|
||||
sha: v.head_sha,
|
||||
state,
|
||||
context: "lakehouse/auditor",
|
||||
description: v.one_liner,
|
||||
target_url: "", // no URL yet; could point to a verdicts dashboard
|
||||
});
|
||||
|
||||
// 2. Issue comment — the reasoning. Gated so we don't spam the PR
|
||||
// with identical comments on re-audits of the same SHA. Caller
|
||||
// (poller) ensures we only re-audit fresh SHAs, but a dedup
|
||||
// marker inside the body keeps it idempotent if re-run.
|
||||
const body = formatReviewBody(v);
|
||||
await postIssueComment({ pr_number: v.pr_number, body });
|
||||
}
|
||||
|
||||
function formatReviewBody(v: Verdict): string {
|
||||
const byCheck: Record<string, Finding[]> = {};
|
||||
for (const f of v.findings) {
|
||||
(byCheck[f.check] ||= []).push(f);
|
||||
}
|
||||
|
||||
const verdictEmoji =
|
||||
v.overall === "approve" ? "✅" :
|
||||
v.overall === "request_changes" ? "⚠️" :
|
||||
"🛑";
|
||||
|
||||
const lines: string[] = [];
|
||||
lines.push(`## Auditor verdict: ${verdictEmoji} \`${v.overall}\``);
|
||||
lines.push("");
|
||||
lines.push(`**One-liner:** ${v.one_liner}`);
|
||||
lines.push(`**Head SHA:** \`${v.head_sha.slice(0, 12)}\``);
|
||||
lines.push(`**Audited at:** ${v.audited_at}`);
|
||||
lines.push("");
|
||||
|
||||
// Per-check sections, only if the check produced findings.
|
||||
const checkOrder = ["static", "dynamic", "inference", "kb_query"] as const;
|
||||
for (const check of checkOrder) {
|
||||
const fs = byCheck[check] ?? [];
|
||||
if (fs.length === 0) continue;
|
||||
const bySev = {
|
||||
block: fs.filter(f => f.severity === "block").length,
|
||||
warn: fs.filter(f => f.severity === "warn").length,
|
||||
info: fs.filter(f => f.severity === "info").length,
|
||||
};
|
||||
lines.push(`<details><summary><b>${check}</b> — ${fs.length} findings (${bySev.block} block, ${bySev.warn} warn, ${bySev.info} info)</summary>`);
|
||||
lines.push("");
|
||||
for (const f of fs) {
|
||||
const mark = f.severity === "block" ? "🛑" : f.severity === "warn" ? "⚠️" : "ℹ️";
|
||||
lines.push(`${mark} **${f.severity}** — ${f.summary}`);
|
||||
for (const e of f.evidence.slice(0, 3)) {
|
||||
lines.push(` - \`${e.slice(0, 180).replace(/\n/g, " ")}\``);
|
||||
}
|
||||
}
|
||||
lines.push("");
|
||||
lines.push("</details>");
|
||||
lines.push("");
|
||||
}
|
||||
|
||||
lines.push("### Metrics");
|
||||
lines.push("```json");
|
||||
lines.push(JSON.stringify(v.metrics, null, 2));
|
||||
lines.push("```");
|
||||
lines.push("");
|
||||
lines.push(`<sub>Lakehouse auditor · SHA ${v.head_sha.slice(0, 8)} · re-audit on new commit flips the status automatically.</sub>`);
|
||||
|
||||
return lines.join("\n");
|
||||
}
|
||||
|
||||
function stubFinding(check: "dynamic" | "inference", why: string): Finding[] {
|
||||
return [{ check, severity: "info", summary: `${check} check skipped — ${why}`, evidence: [why] }];
|
||||
}
|
||||
@@ -1,68 +0,0 @@
|
||||
// One-shot dry-run audit of a single PR. Useful for verifying check
|
||||
// behavior (kb_query scrum surfacing, inference prompts, etc.) without
|
||||
// posting to Gitea. Does NOT touch state.json and does NOT post
|
||||
// commit status or PR comments.
|
||||
//
|
||||
// Run: bun run auditor/audit_one.ts <pr-number>
|
||||
|
||||
import { getPrSnapshot } from "./gitea.ts";
|
||||
import { auditPr } from "./audit.ts";
|
||||
|
||||
async function main() {
|
||||
const prNumRaw = process.argv[2];
|
||||
if (!prNumRaw) {
|
||||
console.error("usage: bun run auditor/audit_one.ts <pr-number>");
|
||||
process.exit(2);
|
||||
}
|
||||
const prNum = Number(prNumRaw);
|
||||
if (!Number.isFinite(prNum)) {
|
||||
console.error(`invalid PR number: ${prNumRaw}`);
|
||||
process.exit(2);
|
||||
}
|
||||
|
||||
console.log(`[audit_one] fetching PR #${prNum}...`);
|
||||
const pr = await getPrSnapshot(prNum);
|
||||
console.log(`[audit_one] PR #${pr.number}: "${pr.title}" (head=${pr.head_sha.slice(0, 12)})`);
|
||||
console.log(`[audit_one] files in diff: ${pr.files.length}`);
|
||||
for (const f of pr.files) console.log(` - ${f.path} (+${f.additions}/-${f.deletions})`);
|
||||
console.log("");
|
||||
|
||||
const verdict = await auditPr(pr, {
|
||||
dry_run: true, // no Gitea posting
|
||||
skip_dynamic: true, // don't run fixture
|
||||
skip_inference: process.env.LH_AUDITOR_SKIP_INFERENCE === "1",
|
||||
});
|
||||
|
||||
console.log("\n═══ VERDICT ═══");
|
||||
console.log(`overall: ${verdict.overall}`);
|
||||
console.log(`one-liner: ${verdict.one_liner}`);
|
||||
console.log(`findings: total=${verdict.metrics.findings_total} block=${verdict.metrics.findings_block} warn=${verdict.metrics.findings_warn} info=${verdict.metrics.findings_info}`);
|
||||
console.log("");
|
||||
|
||||
// Print findings, highlighting kb_query scrum surfacing
|
||||
const byCheck: Record<string, typeof verdict.findings> = {};
|
||||
for (const f of verdict.findings) (byCheck[f.check] ||= []).push(f);
|
||||
|
||||
for (const [check, findings] of Object.entries(byCheck)) {
|
||||
console.log(`── ${check} (${findings.length}) ──`);
|
||||
for (const f of findings) {
|
||||
const tag = f.severity === "block" ? "🛑" : f.severity === "warn" ? "⚠️ " : "ℹ️ ";
|
||||
console.log(` ${tag} [${f.severity}] ${f.summary}`);
|
||||
if (f.summary.includes("scrum-master")) {
|
||||
for (const e of f.evidence) {
|
||||
console.log(` → ${e.slice(0, 200)}`);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const scrumFindings = verdict.findings.filter(f => f.summary.includes("scrum-master"));
|
||||
console.log("");
|
||||
console.log(`═══ SCRUM WIRE CHECK: ${scrumFindings.length} scrum-master findings surfaced by kb_query ═══`);
|
||||
if (scrumFindings.length === 0) {
|
||||
console.log(" (none — either no matching scrum_reviews.jsonl rows, or files didn't match PR diff)");
|
||||
}
|
||||
process.exit(0);
|
||||
}
|
||||
|
||||
main().catch(e => { console.error("[audit_one] fatal:", e); process.exit(1); });
|
||||
@@ -1,91 +0,0 @@
|
||||
// Dynamic execution check — runs the hybrid fixture and maps its
|
||||
// layer results to auditor Findings.
|
||||
//
|
||||
// A layer that fails with a "not implemented / 404 / slice N" error
|
||||
// gets severity=info (honest placeholder signal). A layer that fails
|
||||
// any other way gets severity=warn (something actually broke).
|
||||
// An info-level summary finding is always emitted carrying the real
|
||||
// numbers — shipped/placeholder phase counts, per-layer latency.
|
||||
|
||||
import { runHybridFixture } from "../fixtures/hybrid_38_40_45.ts";
|
||||
import type { Finding } from "../types.ts";
|
||||
|
||||
const PLACEHOLDER_MARKERS = [
|
||||
"unimplemented",
|
||||
" 404 ", "(404)", " 405 ", "(405)",
|
||||
"slice 3", "slice 4", "slice 5",
|
||||
"endpoint not built", "not yet",
|
||||
];
|
||||
|
||||
function isPlaceholderFailure(err?: string): boolean {
|
||||
if (!err) return false;
|
||||
const low = err.toLowerCase();
|
||||
return PLACEHOLDER_MARKERS.some(m => low.includes(m.toLowerCase()));
|
||||
}
|
||||
|
||||
export async function runDynamicCheck(): Promise<Finding[]> {
|
||||
const findings: Finding[] = [];
|
||||
|
||||
let result;
|
||||
try {
|
||||
result = await runHybridFixture();
|
||||
} catch (e) {
|
||||
// Fixture itself crashed — can't run dynamic check at all.
|
||||
return [
|
||||
{
|
||||
check: "dynamic",
|
||||
severity: "warn",
|
||||
summary: `hybrid fixture crashed before completing: ${(e as Error).message.slice(0, 140)}`,
|
||||
evidence: [(e as Error).message],
|
||||
},
|
||||
];
|
||||
}
|
||||
|
||||
// Per-layer findings for every non-ok layer.
|
||||
for (const layer of result.layers) {
|
||||
if (layer.ok) continue;
|
||||
const placeholder = isPlaceholderFailure(layer.error);
|
||||
findings.push({
|
||||
check: "dynamic",
|
||||
severity: placeholder ? "info" : "warn",
|
||||
summary: placeholder
|
||||
? `hybrid fixture layer ${layer.layer} (Phase ${layer.phase}) honestly reports unimplemented`
|
||||
: `hybrid fixture layer ${layer.layer} (Phase ${layer.phase}) failed — not a placeholder, a real failure`,
|
||||
evidence: [
|
||||
`evidence: ${layer.evidence.slice(0, 160)}`,
|
||||
...(layer.error ? [`error: ${layer.error.slice(0, 160)}`] : []),
|
||||
`latency_ms: ${layer.latency_ms}`,
|
||||
],
|
||||
});
|
||||
}
|
||||
|
||||
// One overall summary with real numbers so the report shows what
|
||||
// DID pass plus per-layer timing.
|
||||
const metrics_preview = Object.entries(result.real_numbers)
|
||||
.slice(0, 10)
|
||||
.map(([k, v]) => `${k}=${v}`);
|
||||
findings.push({
|
||||
check: "dynamic",
|
||||
severity: "info",
|
||||
summary: `hybrid fixture overall=${result.overall}, shipped [${result.shipped_phases.join(", ")}], placeholder [${result.placeholder_phases.join(", ")}]`,
|
||||
evidence: metrics_preview.length > 0 ? metrics_preview : ["no metrics emitted"],
|
||||
});
|
||||
|
||||
// If the fixture ran at all but nothing passed, elevate one of the
|
||||
// summary findings to warn — something more than "all honest
|
||||
// placeholders" is wrong.
|
||||
if (result.overall === "fail") {
|
||||
findings.push({
|
||||
check: "dynamic",
|
||||
severity: "warn",
|
||||
summary: `hybrid fixture: 0 layers passed (overall=fail)`,
|
||||
evidence: [
|
||||
"a total fixture fail usually means a precondition service is down",
|
||||
"(gateway /health / sidecar / Langfuse /v1/chat) — NOT necessarily",
|
||||
"the PR's code problem. Check service status before blaming the PR.",
|
||||
],
|
||||
});
|
||||
}
|
||||
|
||||
return findings;
|
||||
}
|
||||
@@ -1,666 +0,0 @@
|
||||
// Cloud inference check — wraps the proven run_codereview pattern
|
||||
// from llm_team_ui.py (same 3-stage framing, same cloud model) to
|
||||
// critique a PR's claims against its diff.
|
||||
//
|
||||
// Proved out 2026-04-22 via /tmp/codereview_runner.py — gpt-oss:120b
|
||||
// caught a real ternary bug in auditor/fixtures/hybrid_38_40_45.ts
|
||||
// that unit tests missed. This module reuses the reviewer prompt
|
||||
// shape (bugs / security / performance / style / edge cases) and
|
||||
// adds claim-vs-diff specific framing.
|
||||
//
|
||||
// Call surface: runInferenceCheck(claims, diff) → Finding[].
|
||||
// Cloud latency budget: ~60s (gpt-oss:120b reviewer typically 35-50s
|
||||
// with a 15KB diff + claim list).
|
||||
|
||||
import type { Claim, Finding } from "../types.ts";
|
||||
import { Glob } from "bun";
|
||||
import { readFile, mkdir, appendFile } from "node:fs/promises";
|
||||
import { extractFacts } from "../fact_extractor.ts";
|
||||
|
||||
const GATEWAY = process.env.LH_GATEWAY_URL ?? "http://localhost:3100";
|
||||
const MODEL = process.env.LH_AUDITOR_REVIEW_MODEL ?? "gpt-oss:120b";
|
||||
// Tie-breaker for claims where the N=3 consensus produces a 1-1-1
|
||||
// split (genuinely borderline). Different architecture from the
|
||||
// primary reviewer (gpt-oss) so the tie-break isn't correlated with
|
||||
// the original disagreement. qwen3-coder:480b is a newer coding
|
||||
// specialist at 480B params, well-suited to PR-diff claim verification
|
||||
// and distinct in training lineage from gpt-oss.
|
||||
const TIEBREAKER_MODEL = process.env.LH_AUDITOR_TIEBREAKER_MODEL ?? "qwen3-coder:480b";
|
||||
const N_CONSENSUS = Number(process.env.LH_AUDITOR_CONSENSUS_N ?? 3);
|
||||
const AUDIT_DISCREPANCIES_JSONL = "/home/profit/lakehouse/data/_kb/audit_discrepancies.jsonl";
|
||||
// 40KB comfortably fits gpt-oss:120b's context. PR #1 (~39KB) was
|
||||
// previously truncated at 15KB causing the reviewer to miss later
|
||||
// files (gitea.ts, policy.ts) and flag "no Gitea client present" as a
|
||||
// block finding when the file was simply outside the truncation window.
|
||||
//
|
||||
// Above this threshold we curate via tree-split rather than truncate,
|
||||
// following the scrum_master pattern: shard the diff, summarize each
|
||||
// shard against the claim-verification task, merge into a compact
|
||||
// scratchpad, then ask the cloud to verify claims against the
|
||||
// scratchpad. This gives the cloud full-PR fidelity without bursting
|
||||
// its context window (observed failure mode: empty response or
|
||||
// unparseable output when prompt exceeds model's comfortable range).
|
||||
const MAX_DIFF_CHARS = 40000;
|
||||
// Tree-split kicks in above this. 30KB is below MAX_DIFF_CHARS so we
|
||||
// curate BEFORE truncation would happen — never lose signal to a hard
|
||||
// cut. Shard size is chosen so ~10 shards cover PR #8-size diffs in a
|
||||
// reasonable round-trip budget.
|
||||
const CURATION_THRESHOLD = 30000;
|
||||
const DIFF_SHARD_SIZE = 4500;
|
||||
const CALL_TIMEOUT_MS = 120_000;
|
||||
const REPO_ROOT = "/home/profit/lakehouse";
|
||||
|
||||
export interface InferenceContext {
|
||||
pr_number: number;
|
||||
head_sha: string;
|
||||
}
|
||||
|
||||
const AUDIT_FACTS_JSONL = "/home/profit/lakehouse/data/_kb/audit_facts.jsonl";
|
||||
|
||||
export async function runInferenceCheck(
|
||||
claims: Claim[],
|
||||
diff: string,
|
||||
ctx?: InferenceContext,
|
||||
): Promise<Finding[]> {
|
||||
if (claims.length === 0) {
|
||||
return [{
|
||||
check: "inference",
|
||||
severity: "info",
|
||||
summary: "no ship-claims extracted — skipping cloud inference",
|
||||
evidence: ["parser returned empty claim list; nothing to verify against cloud"],
|
||||
}];
|
||||
}
|
||||
|
||||
// Empirical claims (runtime metrics / observed outcomes) can't be
|
||||
// verified from the diff. Drop them from the cloud prompt so the
|
||||
// reviewer doesn't chase ghosts. A future `runtime_evidence` check
|
||||
// can validate these against data/_kb/*/summary.json outputs.
|
||||
const verifiable = claims.filter(c => c.strength !== "empirical");
|
||||
const empiricalCount = claims.length - verifiable.length;
|
||||
if (verifiable.length === 0) {
|
||||
return [{
|
||||
check: "inference",
|
||||
severity: "info",
|
||||
summary: `all ${claims.length} claims are empirical (runtime metrics) — skipping cloud inference`,
|
||||
evidence: [`empirical claims can't be verified from a static diff; needs runtime-evidence check`],
|
||||
}];
|
||||
}
|
||||
|
||||
// Diff source for the cloud prompt — either the raw diff (small
|
||||
// enough to fit), or a tree-split scratchpad (curation layer). We
|
||||
// prefer curation to truncation: truncation silently drops files
|
||||
// past the window; curation summarizes them so the cloud still sees
|
||||
// what changed, just densified.
|
||||
let diffForPrompt: string;
|
||||
let curationNote = "";
|
||||
if (diff.length > CURATION_THRESHOLD) {
|
||||
const ts = await treeSplitDiff(diff, verifiable);
|
||||
diffForPrompt = ts.scratchpad;
|
||||
curationNote = ` (curated: ${diff.length} chars → ${ts.shards} shards → scratchpad ${ts.scratchpad.length} chars)`;
|
||||
} else {
|
||||
diffForPrompt = diff;
|
||||
}
|
||||
// Belt-and-suspenders truncation — even a tree-split scratchpad
|
||||
// shouldn't exceed MAX_DIFF_CHARS in practice, but guard anyway so
|
||||
// pathological inputs can't burst the prompt.
|
||||
const truncated = diffForPrompt.length > MAX_DIFF_CHARS
|
||||
? diffForPrompt.slice(0, MAX_DIFF_CHARS) + `\n...[${diffForPrompt.length - MAX_DIFF_CHARS} more chars truncated]`
|
||||
: diffForPrompt;
|
||||
|
||||
// Build the reviewer prompt in the same shape as run_codereview's
|
||||
// review stage (llm_team_ui.py:10950), adapted for claim verification:
|
||||
// "Task: ..."
|
||||
// "Code: ..."
|
||||
// "Review: bugs/security/perf/style/edge. Provide corrected code."
|
||||
// We add: claim list upfront + ask for structured JSON verdict.
|
||||
//
|
||||
// When the diff was curated (tree-split scratchpad), we add an
|
||||
// explicit anti-false-positive instruction: the scratchpad is a
|
||||
// distillation, not the full source, so absence-from-scratchpad is
|
||||
// NOT evidence of absence-from-diff. Mirrors the fix we made in
|
||||
// scrum_master's review prompt for the same class of error.
|
||||
const isCurated = curationNote.length > 0;
|
||||
const curationGuard = isCurated
|
||||
? [
|
||||
"",
|
||||
"CRITICAL: the 'Diff' below is a curated multi-shard scratchpad,",
|
||||
"NOT the full raw diff. The scratchpad distills each shard down",
|
||||
"to facts useful for claim verification and drops the rest.",
|
||||
"DO NOT flag a function/field/feature as 'missing' or 'not",
|
||||
"implemented' based solely on its absence from the scratchpad —",
|
||||
"absence in a distillation is NOT evidence of absence in the",
|
||||
"actual diff. Only judge a claim NOT BACKED when the scratchpad",
|
||||
"DIRECTLY contradicts it (e.g. scratchpad shows the function was",
|
||||
"added empty, or shows the claimed code path is a stub).",
|
||||
"Skip the unflagged_gaps section entirely when operating on a",
|
||||
"curated scratchpad — you can't reliably detect gaps from a",
|
||||
"distillation, and false positives there are worse than misses.",
|
||||
].join("\n")
|
||||
: "";
|
||||
const systemMsg = [
|
||||
"You review pull-request diffs against the author's own ship-claims.",
|
||||
"For each claim, decide: is it backed by actual code in the diff, or is",
|
||||
"it placeholder / aspirational / unwired?",
|
||||
"",
|
||||
"A claim is BACKED when the diff contains a real code path that delivers",
|
||||
"the claimed behavior. A claim is NOT BACKED when:",
|
||||
" - the claim asserts functionality but the diff only adds types/fields",
|
||||
" with no consumer",
|
||||
" - the claim mentions tests but no test function was added",
|
||||
" - the claim claims integration but the integration point is a stub",
|
||||
" - the diff contains unimplemented!() / todo!() / TODO comments",
|
||||
" - the claim says 'works end-to-end' but the diff has no end-to-end test",
|
||||
curationGuard,
|
||||
"",
|
||||
"Respond with strict JSON only. No prose before or after. Shape:",
|
||||
"{",
|
||||
' "claim_verdicts": [',
|
||||
' {"claim_idx": 0, "backed": false, "evidence": "short reason"}',
|
||||
" ],",
|
||||
' "unflagged_gaps": [',
|
||||
' {"location": "file:line", "summary": "short description"}',
|
||||
" ]",
|
||||
"}",
|
||||
].join("\n");
|
||||
|
||||
const userMsg = [
|
||||
`Ship-claims the author made (numbered 0..N-1):`,
|
||||
verifiable.map((c, i) => ` ${i}. [${c.strength}] "${c.text}" at ${c.location}`).join("\n"),
|
||||
"",
|
||||
`Diff:`,
|
||||
"```",
|
||||
truncated,
|
||||
"```",
|
||||
"",
|
||||
`For each numbered claim above, emit a claim_verdicts entry. For gaps the`,
|
||||
`author DIDN'T claim but that look like placeholder code, emit unflagged_gaps.`,
|
||||
`Strict JSON only, matching the shape described. No prose outside JSON.`,
|
||||
].join("\n");
|
||||
|
||||
// N=3 consensus — run the primary reviewer in parallel, collect
|
||||
// all three parsed responses, majority-vote per claim. Parallel
|
||||
// (Promise.all) because each call is ~20-30s and they're independent;
|
||||
// wall-clock stays ~same as single call, cost 3x tokens. Empirical
|
||||
// justification: in 3-run determinism tests, 7/8 findings were
|
||||
// stable but 1 flipped across runs — majority vote stabilizes the
|
||||
// flipping class without losing the stable signal.
|
||||
const primaryRuns = await Promise.all(
|
||||
Array.from({ length: N_CONSENSUS }, () =>
|
||||
runCloudInference(systemMsg, userMsg, MODEL)),
|
||||
);
|
||||
|
||||
const parsedRuns = primaryRuns.filter(r => r.parsed !== null);
|
||||
if (parsedRuns.length === 0) {
|
||||
// All N calls failed. Surface the first-run diagnostic so the
|
||||
// operator sees *why* (unreachable / non-200 / unparseable).
|
||||
const first = primaryRuns[0];
|
||||
return [{
|
||||
check: "inference",
|
||||
severity: "info",
|
||||
summary: `cloud inference all ${N_CONSENSUS} consensus runs failed — ${first.error ?? "unknown"}`,
|
||||
evidence: [
|
||||
`first-run diagnostic: ${first.diagnostic ?? "(none)"}`,
|
||||
`successful runs: 0 / ${N_CONSENSUS}`,
|
||||
],
|
||||
}];
|
||||
}
|
||||
|
||||
// Aggregate votes per claim_idx.
|
||||
interface Votes { trues: number; falses: number; evidences: string[] }
|
||||
const votesByClaim = new Map<number, Votes>();
|
||||
const unflaggedByRun: any[][] = [];
|
||||
let totalTokens = 0;
|
||||
for (const run of parsedRuns) {
|
||||
totalTokens += run.tokens;
|
||||
unflaggedByRun.push(Array.isArray(run.parsed?.unflagged_gaps) ? run.parsed.unflagged_gaps : []);
|
||||
for (const v of run.parsed?.claim_verdicts ?? []) {
|
||||
const idx = Number(v?.claim_idx);
|
||||
if (!Number.isFinite(idx)) continue;
|
||||
const rec = votesByClaim.get(idx) ?? { trues: 0, falses: 0, evidences: [] };
|
||||
if (v.backed === false) {
|
||||
rec.falses++;
|
||||
rec.evidences.push(String(v.evidence ?? ""));
|
||||
} else if (v.backed === true) {
|
||||
rec.trues++;
|
||||
}
|
||||
votesByClaim.set(idx, rec);
|
||||
}
|
||||
}
|
||||
|
||||
const findings: Finding[] = [];
|
||||
|
||||
// Summary finding so the verdict layer knows the check ran.
|
||||
findings.push({
|
||||
check: "inference",
|
||||
severity: "info",
|
||||
summary: `cloud review completed (model=${MODEL}, consensus=${parsedRuns.length}/${N_CONSENSUS}, tokens=${totalTokens})${curationNote}`,
|
||||
evidence: [
|
||||
`claims voted: ${votesByClaim.size}`,
|
||||
`parsed runs: ${parsedRuns.length} / ${N_CONSENSUS}`,
|
||||
],
|
||||
});
|
||||
|
||||
// Per-claim majority vote; tie-break if no majority.
|
||||
const discrepancies: Array<{
|
||||
claim_idx: number;
|
||||
claim_text: string;
|
||||
votes: { trues: number; falses: number };
|
||||
resolution: "majority_backed" | "majority_not_backed" | "tiebreaker_backed" | "tiebreaker_not_backed" | "unresolved";
|
||||
tiebreaker_model?: string;
|
||||
}> = [];
|
||||
|
||||
for (const [idx, votes] of votesByClaim) {
|
||||
const claim = verifiable[idx];
|
||||
if (!claim) continue;
|
||||
const totalVotes = votes.trues + votes.falses;
|
||||
let notBacked: boolean | null = null;
|
||||
let resolution: typeof discrepancies[number]["resolution"] = "majority_backed";
|
||||
let evidenceText = "";
|
||||
let tbModel: string | undefined;
|
||||
|
||||
if (votes.falses > votes.trues) {
|
||||
notBacked = true;
|
||||
resolution = "majority_not_backed";
|
||||
evidenceText = votes.evidences[0] ?? "(no reason given)";
|
||||
} else if (votes.trues > votes.falses) {
|
||||
notBacked = false;
|
||||
resolution = "majority_backed";
|
||||
} else {
|
||||
// Tie. Run tie-breaker with a different-architecture model.
|
||||
const tb = await runCloudInference(systemMsg, userMsg, TIEBREAKER_MODEL);
|
||||
if (tb.parsed) {
|
||||
const tv = (tb.parsed.claim_verdicts ?? []).find((v: any) => Number(v?.claim_idx) === idx);
|
||||
if (tv?.backed === false) {
|
||||
notBacked = true;
|
||||
resolution = "tiebreaker_not_backed";
|
||||
evidenceText = `(tie-breaker ${TIEBREAKER_MODEL}) ${String(tv.evidence ?? "")}`;
|
||||
tbModel = TIEBREAKER_MODEL;
|
||||
} else if (tv?.backed === true) {
|
||||
notBacked = false;
|
||||
resolution = "tiebreaker_backed";
|
||||
tbModel = TIEBREAKER_MODEL;
|
||||
} else {
|
||||
resolution = "unresolved";
|
||||
}
|
||||
} else {
|
||||
resolution = "unresolved";
|
||||
}
|
||||
}
|
||||
|
||||
// Log every case where the N runs disagreed — discrepancies are
|
||||
// signal, not noise. Separate from audit_lessons.jsonl because
|
||||
// they're about the *auditor's* quality, not the PR's quality.
|
||||
const disagreed = totalVotes >= 2 && votes.trues > 0 && votes.falses > 0;
|
||||
if (disagreed || resolution.startsWith("tiebreaker") || resolution === "unresolved") {
|
||||
discrepancies.push({
|
||||
claim_idx: idx,
|
||||
claim_text: claim.text,
|
||||
votes: { trues: votes.trues, falses: votes.falses },
|
||||
resolution,
|
||||
tiebreaker_model: tbModel,
|
||||
});
|
||||
}
|
||||
|
||||
if (notBacked === true) {
|
||||
const sev: Finding["severity"] = claim.strength === "strong" ? "block"
|
||||
: claim.strength === "moderate" ? "warn"
|
||||
: "info";
|
||||
findings.push({
|
||||
check: "inference",
|
||||
severity: sev,
|
||||
claim_text: claim.text,
|
||||
summary: `cloud: claim not backed — "${claim.text.slice(0, 100)}"`,
|
||||
evidence: [
|
||||
`at ${claim.location}`,
|
||||
`consensus: ${votes.falses}/${totalVotes} not-backed (resolution: ${resolution})`,
|
||||
`cloud reason: ${evidenceText.slice(0, 200)}`,
|
||||
],
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Persist discrepancies so we can measure consensus drift over time.
|
||||
if (discrepancies.length > 0 && ctx) {
|
||||
persistDiscrepancies(ctx, discrepancies).catch(e =>
|
||||
console.error(`[inference] discrepancy log failed: ${(e as Error).message}`));
|
||||
}
|
||||
|
||||
// Use first run's parsed for downstream unflagged_gaps processing.
|
||||
const parsed = parsedRuns[0].parsed;
|
||||
|
||||
// Route the curated scratchpad through llm_team's extract-facts
|
||||
// pipeline when we have (a) a curated scratchpad (best signal about
|
||||
// what the PR actually changed) and (b) PR context to scope facts.
|
||||
// AWAITED (not fire-and-forget) so CLI callers like audit_one.ts
|
||||
// don't exit before extraction lands; the systemd poller has plenty
|
||||
// of headroom (90s cycle vs ~15s extraction). A failure inside
|
||||
// extractAndPersistFacts is caught + logged but never throws.
|
||||
if (isCurated && ctx && process.env.LH_AUDITOR_SKIP_EXTRACT !== "1") {
|
||||
try {
|
||||
await extractAndPersistFacts(diffForPrompt, ctx);
|
||||
} catch (e) {
|
||||
console.error(`[inference] fact extraction failed: ${(e as Error).message}`);
|
||||
}
|
||||
}
|
||||
|
||||
// Belt-and-suspenders: when operating on a curated scratchpad, drop
|
||||
// the unflagged_gaps section entirely. The distillation can't
|
||||
// reliably ground gap-detection, and false positives are worse than
|
||||
// misses for this signal class. The systemMsg already asks the
|
||||
// cloud to skip this section when curated — but the model may still
|
||||
// emit it, so we filter here too.
|
||||
const gapsToEmit = isCurated ? [] : (parsed.unflagged_gaps ?? []);
|
||||
for (const g of gapsToEmit) {
|
||||
const summary = String(g?.summary ?? "?");
|
||||
const location = String(g?.location ?? "?");
|
||||
// False-positive guard — when the cloud says "X not defined in this
|
||||
// diff" or "missing implementation of X", the cloud may just mean
|
||||
// "X is not in the added lines," not "X doesn't exist in the repo."
|
||||
// Extract candidate symbol names and grep the repo. If any symbol
|
||||
// is defined elsewhere, drop the finding — it's a known-symbol
|
||||
// reference, not a placeholder.
|
||||
if (/not\s+defined|missing\s+implementation|never\s+referenced\s+or\s+integrated/i.test(summary)) {
|
||||
const symbols = extractSymbols(summary);
|
||||
if (symbols.length > 0) {
|
||||
const resolved = await symbolsExistInRepo(symbols);
|
||||
if (resolved.length === symbols.length) {
|
||||
// Every named symbol exists somewhere in the repo — silent drop.
|
||||
continue;
|
||||
}
|
||||
if (resolved.length > 0) {
|
||||
// Partially resolved — demote to info with a note.
|
||||
findings.push({
|
||||
check: "inference",
|
||||
severity: "info",
|
||||
summary: `cloud gap partially resolved by repo grep: ${summary.slice(0, 120)}`,
|
||||
evidence: [
|
||||
`location: ${location.slice(0, 140)}`,
|
||||
`resolved via grep: ${resolved.join(",")}`,
|
||||
`unresolved: ${symbols.filter(s => !resolved.includes(s)).join(",")}`,
|
||||
],
|
||||
});
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
findings.push({
|
||||
check: "inference",
|
||||
severity: "warn",
|
||||
summary: `cloud-flagged gap not in any claim: ${summary.slice(0, 120)}`,
|
||||
evidence: [`location: ${location.slice(0, 140)}`],
|
||||
});
|
||||
}
|
||||
|
||||
return findings;
|
||||
}
|
||||
|
||||
// Single cloud call — the consensus loop calls this N times in
|
||||
// parallel. Returns the parsed JSON shape + token usage + any error
|
||||
// diagnostic. NEVER throws; the consensus aggregator handles partial
|
||||
// failures by dropping non-parsed runs from the vote.
|
||||
interface CloudRunResult {
|
||||
parsed: any | null;
|
||||
tokens: number;
|
||||
error?: string; // "unreachable" | "non_200" | "unparseable"
|
||||
diagnostic?: string; // first 200 chars for debugging
|
||||
model: string;
|
||||
}
|
||||
|
||||
async function runCloudInference(systemMsg: string, userMsg: string, model: string): Promise<CloudRunResult> {
|
||||
let resp: Response;
|
||||
try {
|
||||
resp = await fetch(`${GATEWAY}/v1/chat`, {
|
||||
method: "POST",
|
||||
headers: { "content-type": "application/json" },
|
||||
body: JSON.stringify({
|
||||
provider: "ollama_cloud",
|
||||
model,
|
||||
messages: [
|
||||
{ role: "system", content: systemMsg },
|
||||
{ role: "user", content: userMsg },
|
||||
],
|
||||
// temp=0 (greedy) + think=true. think=true is required for
|
||||
// gpt-oss:120b — without it the model returns empty content
|
||||
// on large prompts. Variance from the think trace is observed
|
||||
// in practice, which is why we use N=3 consensus, not single-
|
||||
// call determinism.
|
||||
max_tokens: 3000,
|
||||
temperature: 0,
|
||||
think: true,
|
||||
}),
|
||||
signal: AbortSignal.timeout(CALL_TIMEOUT_MS),
|
||||
});
|
||||
} catch (e) {
|
||||
return { parsed: null, tokens: 0, error: "unreachable", diagnostic: (e as Error).message.slice(0, 200), model };
|
||||
}
|
||||
if (!resp.ok) {
|
||||
return { parsed: null, tokens: 0, error: "non_200", diagnostic: `${resp.status}: ${(await resp.text()).slice(0, 160)}`, model };
|
||||
}
|
||||
let body: any;
|
||||
try { body = await resp.json(); }
|
||||
catch (e) { return { parsed: null, tokens: 0, error: "unparseable", diagnostic: (e as Error).message, model }; }
|
||||
const content: string = body?.choices?.[0]?.message?.content ?? "";
|
||||
const tokens: number = body?.usage?.total_tokens ?? 0;
|
||||
const parsed = extractJson(content);
|
||||
if (!parsed) {
|
||||
return { parsed: null, tokens, error: "unparseable", diagnostic: content.slice(0, 200), model };
|
||||
}
|
||||
return { parsed, tokens, model };
|
||||
}
|
||||
|
||||
async function persistDiscrepancies(ctx: InferenceContext, discrepancies: any[]): Promise<void> {
|
||||
await mkdir("/home/profit/lakehouse/data/_kb", { recursive: true });
|
||||
const rows = discrepancies.map(d => JSON.stringify({
|
||||
pr_number: ctx.pr_number,
|
||||
head_sha: ctx.head_sha,
|
||||
logged_at: new Date().toISOString(),
|
||||
...d,
|
||||
}));
|
||||
await appendFile(AUDIT_DISCREPANCIES_JSONL, rows.join("\n") + "\n");
|
||||
}
|
||||
|
||||
// Extract structured knowledge from the curated scratchpad and append
|
||||
// to data/_kb/audit_facts.jsonl — one row per extract run, keyed by
|
||||
// PR number + head SHA for scope tracking. kb_query tails this next
|
||||
// audit to surface recurring entities/relationships across PRs.
|
||||
async function extractAndPersistFacts(scratchpad: string, ctx: InferenceContext): Promise<void> {
|
||||
const ex = await extractFacts(scratchpad);
|
||||
if (ex.error && ex.entities.length === 0 && ex.facts.length === 0) {
|
||||
// Full failure — log but don't write an empty row.
|
||||
console.error(`[inference] extractFacts skipped row: ${ex.error}`);
|
||||
return;
|
||||
}
|
||||
const row = {
|
||||
pr_number: ctx.pr_number,
|
||||
head_sha: ctx.head_sha,
|
||||
extracted_at: ex.extracted_at,
|
||||
extractor: ex.extractor_model,
|
||||
verifier: ex.verifier_model,
|
||||
llm_team_run_id: ex.llm_team_run_id ?? null,
|
||||
facts: ex.facts,
|
||||
entities: ex.entities,
|
||||
relationships: ex.relationships,
|
||||
verification_preview: ex.verification.slice(0, 400),
|
||||
verifier_verdicts: ex.verifier_verdicts,
|
||||
facts_dropped_by_verifier: ex.facts_dropped_by_verifier ?? 0,
|
||||
schema_version: 2,
|
||||
source: "audit_inference",
|
||||
};
|
||||
await mkdir("/home/profit/lakehouse/data/_kb", { recursive: true });
|
||||
await appendFile(AUDIT_FACTS_JSONL, JSON.stringify(row) + "\n");
|
||||
}
|
||||
|
||||
// Curation via tree-split — ports the scrum_master pattern into the
|
||||
// inference check. Shards the raw diff into DIFF_SHARD_SIZE chunks,
|
||||
// summarizes each shard *against the claim-verification task* so the
|
||||
// summary preserves exactly what the cloud needs to judge claims
|
||||
// (function signatures, struct fields, deletions, new files), drops
|
||||
// everything else. Merges into a compact scratchpad.
|
||||
//
|
||||
// Cost: N cloud calls for the shard summaries + 1 cloud call for the
|
||||
// final verification = N+1 calls instead of 1. Mitigation: shards run
|
||||
// serially (not parallel) to keep gateway load bounded; summary calls
|
||||
// use max_tokens=400 so they're fast (~2s each on gpt-oss:120b).
|
||||
//
|
||||
// Determinism: each shard summary call uses temp=0 + think=true (same
|
||||
// as the top-level inference call), so identical input yields
|
||||
// identical scratchpad. The final verification call then sees a
|
||||
// stable scratchpad, giving stable verdicts.
|
||||
async function treeSplitDiff(
|
||||
fullDiff: string,
|
||||
claims: Claim[],
|
||||
): Promise<{ scratchpad: string; shards: number }> {
|
||||
const shards: Array<{ from: number; to: number; text: string }> = [];
|
||||
for (let i = 0; i < fullDiff.length; i += DIFF_SHARD_SIZE) {
|
||||
const end = Math.min(i + DIFF_SHARD_SIZE, fullDiff.length);
|
||||
shards.push({ from: i, to: end, text: fullDiff.slice(i, end) });
|
||||
}
|
||||
// Curate the claim list into a short form the summary prompt can
|
||||
// use to bias extraction toward relevant facts.
|
||||
const claimDigest = claims.map((c, i) =>
|
||||
`${i}. [${c.strength}] "${c.text.slice(0, 100)}"`
|
||||
).join("\n");
|
||||
|
||||
let scratchpad = "";
|
||||
for (const [si, shard] of shards.entries()) {
|
||||
const prompt = [
|
||||
`You are summarizing shard ${si + 1}/${shards.length} (chars ${shard.from}..${shard.to}) of a PR diff.`,
|
||||
`The downstream task will verify these ship-claims against the full-PR summary. Extract ONLY facts that could confirm or refute these claims:`,
|
||||
"",
|
||||
claimDigest,
|
||||
"",
|
||||
"Extract: new function/method signatures, struct fields, deletions, new files, wiring (function X calls Y), absence-of-implementation markers, TODO comments on added lines.",
|
||||
"Skip: comment-only edits, whitespace, import reordering, unrelated cosmetic changes.",
|
||||
"",
|
||||
"─────── shard diff ───────",
|
||||
shard.text,
|
||||
"─────── end shard ───────",
|
||||
"",
|
||||
"Output: up to 180 words of facts in bullet form. No prose preamble, no claim verdicts (that's for the downstream step).",
|
||||
].join("\n");
|
||||
|
||||
const r = await callCloud(prompt, 400);
|
||||
if (r.content) {
|
||||
scratchpad += `\n--- shard ${si + 1} (chars ${shard.from}..${shard.to}) ---\n${r.content.trim()}\n`;
|
||||
}
|
||||
}
|
||||
return { scratchpad: scratchpad.trim(), shards: shards.length };
|
||||
}
|
||||
|
||||
// Minimal cloud caller used only by treeSplitDiff — same gateway +
|
||||
// model as the top-level call, but think=false. Shards are small
|
||||
// (≤DIFF_SHARD_SIZE ~4500 chars) and the task is pure fact
|
||||
// extraction, not reasoning. think=true on the shards introduced
|
||||
// variance in reasoning traces that compounded across 23 calls into
|
||||
// a non-deterministic scratchpad (observed during curation
|
||||
// validation: same-SHA runs produced 5/7/8 final findings).
|
||||
// think=false on small prompts is stable — only breaks at the main
|
||||
// call's 10K+ prompt size, which keeps think=true.
|
||||
async function callCloud(prompt: string, maxTokens: number): Promise<{ content: string }> {
|
||||
try {
|
||||
const r = await fetch(`${GATEWAY}/v1/chat`, {
|
||||
method: "POST",
|
||||
headers: { "content-type": "application/json" },
|
||||
body: JSON.stringify({
|
||||
provider: "ollama_cloud",
|
||||
model: MODEL,
|
||||
messages: [{ role: "user", content: prompt }],
|
||||
max_tokens: maxTokens,
|
||||
temperature: 0,
|
||||
think: false,
|
||||
}),
|
||||
signal: AbortSignal.timeout(CALL_TIMEOUT_MS),
|
||||
});
|
||||
if (!r.ok) return { content: "" };
|
||||
const j: any = await r.json();
|
||||
return { content: j?.choices?.[0]?.message?.content ?? "" };
|
||||
} catch {
|
||||
return { content: "" };
|
||||
}
|
||||
}
|
||||
|
||||
// Pull out plausible code-symbol names from a summary string.
|
||||
// Matches:
|
||||
// - identifier with backticks: `foo_bar`
|
||||
// - identifier followed by parens: foo_bar()
|
||||
// - CamelCase types
|
||||
// - snake_case_functions
|
||||
// Filters out common English words that could be matched accidentally.
|
||||
const STOPWORDS = new Set([
|
||||
"not","the","and","for","this","that","with","but","are","was","has",
|
||||
"have","been","any","missing","implementation","diff","defined","never",
|
||||
"referenced","integrated","flow","code","file","some","only","when",
|
||||
]);
|
||||
function extractSymbols(text: string): string[] {
|
||||
const out = new Set<string>();
|
||||
// `backticked` symbols
|
||||
for (const m of text.matchAll(/`([A-Za-z_][A-Za-z0-9_]{2,})`/g)) out.add(m[1]);
|
||||
// foo() or foo_bar() calls
|
||||
for (const m of text.matchAll(/\b([A-Za-z_][A-Za-z0-9_]{2,})\s*\(/g)) out.add(m[1]);
|
||||
// CamelCase types (3+ chars, must start with uppercase)
|
||||
for (const m of text.matchAll(/\b([A-Z][A-Za-z0-9]{2,})\b/g)) out.add(m[1]);
|
||||
return Array.from(out).filter(s => !STOPWORDS.has(s.toLowerCase()));
|
||||
}
|
||||
|
||||
// Scan the repo for at least one definition of each symbol. Uses Bun's
|
||||
// Glob to walk TS/Rust/Python/JS sources; ignores node_modules, data/,
|
||||
// and target/. Skips files > 500KB — those are fixtures/snapshots that
|
||||
// won't contain a definition line and slurping them slows the audit.
|
||||
async function symbolsExistInRepo(symbols: string[]): Promise<string[]> {
|
||||
const patterns = ["**/*.ts", "**/*.tsx", "**/*.rs", "**/*.py", "**/*.js"];
|
||||
const skip = (p: string) => p.includes("/node_modules/") || p.startsWith("data/") || p.includes("/target/") || p.startsWith("dist/");
|
||||
const MAX_FILE_BYTES = 500_000;
|
||||
const { stat } = await import("node:fs/promises");
|
||||
const resolved = new Set<string>();
|
||||
const toFind = new Set(symbols);
|
||||
for (const pat of patterns) {
|
||||
if (toFind.size === 0) break;
|
||||
const glob = new Glob(pat);
|
||||
for await (const f of glob.scan({ cwd: REPO_ROOT, onlyFiles: true })) {
|
||||
if (skip(f)) continue;
|
||||
try { const s = await stat(`${REPO_ROOT}/${f}`); if (s.size > MAX_FILE_BYTES) continue; } catch { continue; }
|
||||
let content: string;
|
||||
try { content = await readFile(`${REPO_ROOT}/${f}`, "utf8"); } catch { continue; }
|
||||
for (const sym of Array.from(toFind)) {
|
||||
// Definition heuristics: `function sym`, `fn sym`, `const sym`,
|
||||
// `let sym`, `def sym`, `class sym`, `struct sym`, `enum sym`,
|
||||
// `trait sym`, `async function sym`, `pub (async )?fn sym`.
|
||||
const re = new RegExp(
|
||||
`\\b(function|async\\s+function|const|let|var|def|class|struct|enum|trait|impl|type|interface|fn|pub\\s+(async\\s+)?fn)\\s+${escapeRe(sym)}\\b`
|
||||
);
|
||||
if (re.test(content)) {
|
||||
resolved.add(sym);
|
||||
toFind.delete(sym);
|
||||
if (toFind.size === 0) break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return Array.from(resolved);
|
||||
}
|
||||
|
||||
function escapeRe(s: string): string {
|
||||
return s.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
|
||||
}
|
||||
|
||||
// Lift the first balanced JSON object out of the response. Tolerates
|
||||
// leading prose, code fences, and model reasoning preamble when the
|
||||
// cloud model ignored "strict JSON only."
|
||||
function extractJson(text: string): any | null {
|
||||
const cleaned = text.replace(/^```(?:json)?\s*/im, "").replace(/```\s*$/im, "");
|
||||
let depth = 0;
|
||||
let start = -1;
|
||||
for (let i = 0; i < cleaned.length; i++) {
|
||||
const c = cleaned[i];
|
||||
if (c === "{") {
|
||||
if (depth === 0) start = i;
|
||||
depth++;
|
||||
} else if (c === "}") {
|
||||
depth--;
|
||||
if (depth === 0 && start >= 0) {
|
||||
try { return JSON.parse(cleaned.slice(start, i + 1)); } catch { start = -1; }
|
||||
}
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
@@ -1,397 +0,0 @@
|
||||
// Local-KB check — reads data/_kb/ + data/_observer/ + data/_bot/
|
||||
// for prior evidence bearing on this PR's claims. Cheap, offline,
|
||||
// no model calls. The point: if a claim like "Phase X shipped" has
|
||||
// a historical record of failing on the same signature before, the
|
||||
// auditor surfaces that pattern before the cloud check has to
|
||||
// infer it.
|
||||
//
|
||||
// What this check reads (all file-backed, append-only or periodic):
|
||||
// data/_kb/outcomes.jsonl — per-scenario outcomes (kb.ts)
|
||||
// data/_kb/error_corrections.jsonl — fail→succeed deltas on same sig
|
||||
// data/_kb/scrum_reviews.jsonl — scrum-master accepted reviews
|
||||
// data/_observer/ops.jsonl — observer ring → disk stream
|
||||
// data/_bot/cycles/*.json — bot cycle results
|
||||
//
|
||||
// Each JSONL line / per-cycle file is small; this check reads tails
|
||||
// only (last N lines or last M files) to stay cheap on large corpora.
|
||||
|
||||
import { readFile, readdir, stat } from "node:fs/promises";
|
||||
import { join } from "node:path";
|
||||
import type { Claim, Finding } from "../types.ts";
|
||||
import { aggregate, ratingSeverity, formatAgg } from "../kb_index.ts";
|
||||
|
||||
const KB_DIR = "/home/profit/lakehouse/data/_kb";
|
||||
const OBSERVER_OPS = "/home/profit/lakehouse/data/_observer/ops.jsonl";
|
||||
const BOT_CYCLES_DIR = "/home/profit/lakehouse/data/_bot/cycles";
|
||||
const SCRUM_REVIEWS_JSONL = "/home/profit/lakehouse/data/_kb/scrum_reviews.jsonl";
|
||||
const AUDIT_LESSONS_JSONL = "/home/profit/lakehouse/data/_kb/audit_lessons.jsonl";
|
||||
const AUDIT_FACTS_JSONL = "/home/profit/lakehouse/data/_kb/audit_facts.jsonl";
|
||||
const TAIL_LINES = 500;
|
||||
const MAX_BOT_CYCLE_FILES = 30;
|
||||
|
||||
export async function runKbCheck(claims: Claim[], prFiles: string[] = []): Promise<Finding[]> {
|
||||
const findings: Finding[] = [];
|
||||
|
||||
// 1. Recent scenario outcomes: are strong-claim-style phrases showing
|
||||
// up alongside failed events? That's "we claimed it worked" +
|
||||
// "it didn't" in the KB.
|
||||
const scenarioFindings = await checkScenarioOutcomes(claims);
|
||||
findings.push(...scenarioFindings);
|
||||
|
||||
// 2. Error corrections: any of the claims text overlap a
|
||||
// recently-observed fail→succeed pair? If yes, add context.
|
||||
const correctionFindings = await checkErrorCorrections(claims);
|
||||
findings.push(...correctionFindings);
|
||||
|
||||
// 3. Bot cycles: any prior bot cycle ended in tests_failed or
|
||||
// apply_failed on a file this PR is also touching?
|
||||
const botFindings = await checkBotCycles();
|
||||
findings.push(...botFindings);
|
||||
|
||||
// 4. Observer: count recent error events. High volume = shared
|
||||
// infra problem, worth flagging (context for other findings).
|
||||
const obsFindings = await checkObserverStream();
|
||||
findings.push(...obsFindings);
|
||||
|
||||
// 5. Scrum-master reviews — surface prior accepted reviews for any
|
||||
// file in this PR's diff. Cohesion plan Phase C wire: the
|
||||
// auditor gets to "borrow" the scrum-master's deeper per-file
|
||||
// analysis instead of re-doing that work.
|
||||
if (prFiles.length > 0) {
|
||||
const scrumFindings = await checkScrumReviews(prFiles);
|
||||
findings.push(...scrumFindings);
|
||||
}
|
||||
|
||||
// 6b. Audit-facts (llm_team extract pipeline output) — surface
|
||||
// entities that recur across multiple PRs. These are the
|
||||
// "core system entities" accumulating in the knowledge base;
|
||||
// showing them as info on future audits gives reviewers
|
||||
// architectural context the raw diff doesn't convey.
|
||||
const factFindings = await checkAuditFacts();
|
||||
findings.push(...factFindings);
|
||||
|
||||
// 6. Audit-lessons feedback loop — summarize the top recurring
|
||||
// patterns from prior audits' block/warn findings. If the same
|
||||
// pattern signature has fired 3+ times across prior audits,
|
||||
// emit it as a block-severity finding so reviewers know this
|
||||
// is a known-recurring class, not a one-off. Does NOT couple
|
||||
// to the current audit's static/inference findings (those run
|
||||
// in parallel and we can't see them here) — the amplification
|
||||
// is emergent: if the current audit's finding-summary matches
|
||||
// a top recurrence, the reviewer sees both.
|
||||
const auditLessonFindings = await checkAuditLessons();
|
||||
findings.push(...auditLessonFindings);
|
||||
|
||||
return findings;
|
||||
}
|
||||
|
||||
async function tailJsonl<T = any>(path: string, n: number): Promise<T[]> {
|
||||
try {
|
||||
const raw = await readFile(path, "utf8");
|
||||
const lines = raw.split("\n").filter(l => l.length > 0);
|
||||
const slice = lines.slice(-n);
|
||||
const out: T[] = [];
|
||||
for (const line of slice) {
|
||||
try { out.push(JSON.parse(line)); } catch { /* skip malformed */ }
|
||||
}
|
||||
return out;
|
||||
} catch {
|
||||
return [];
|
||||
}
|
||||
}
|
||||
|
||||
async function checkScenarioOutcomes(_claims: Claim[]): Promise<Finding[]> {
|
||||
const outcomes = await tailJsonl<any>(join(KB_DIR, "outcomes.jsonl"), TAIL_LINES);
|
||||
if (outcomes.length === 0) return [];
|
||||
const totalEvents = outcomes.reduce((s, o) => s + (o.total_events ?? 0), 0);
|
||||
const okEvents = outcomes.reduce((s, o) => s + (o.ok_events ?? 0), 0);
|
||||
const failRate = totalEvents > 0 ? 1 - okEvents / totalEvents : 0;
|
||||
|
||||
if (totalEvents === 0) {
|
||||
return [{
|
||||
check: "kb_query",
|
||||
severity: "info",
|
||||
summary: `KB: no scenario outcomes on file — learning loop is empty`,
|
||||
evidence: [`data/_kb/outcomes.jsonl has ${outcomes.length} entries with 0 total events`],
|
||||
}];
|
||||
}
|
||||
|
||||
const recent = outcomes.slice(-10);
|
||||
const recentFailSigs: string[] = recent
|
||||
.filter(o => (o.ok_events ?? 0) < (o.total_events ?? 0))
|
||||
.map(o => o.sig_hash)
|
||||
.filter(s => typeof s === "string");
|
||||
|
||||
const findings: Finding[] = [{
|
||||
check: "kb_query",
|
||||
severity: failRate > 0.3 ? "warn" : "info",
|
||||
summary: `KB: ${outcomes.length} recent scenario runs, ${okEvents}/${totalEvents} events ok (fail rate ${(failRate * 100).toFixed(1)}%)`,
|
||||
evidence: [
|
||||
`most recent: ${recent[recent.length - 1]?.run_id ?? "?"}`,
|
||||
`recent failing sigs: ${recentFailSigs.length > 0 ? recentFailSigs.slice(-3).join(", ") : "none"}`,
|
||||
],
|
||||
}];
|
||||
return findings;
|
||||
}
|
||||
|
||||
async function checkErrorCorrections(_claims: Claim[]): Promise<Finding[]> {
|
||||
const corrections = await tailJsonl<any>(join(KB_DIR, "error_corrections.jsonl"), TAIL_LINES);
|
||||
if (corrections.length === 0) return [];
|
||||
return [{
|
||||
check: "kb_query",
|
||||
severity: "info",
|
||||
summary: `KB: ${corrections.length} error corrections on file (fail→succeed pairs)`,
|
||||
evidence: [
|
||||
corrections.length > 0
|
||||
? `most recent: ${String(corrections[corrections.length - 1]?.sig_hash ?? "?").slice(0, 24)}`
|
||||
: "none",
|
||||
],
|
||||
}];
|
||||
}
|
||||
|
||||
async function checkBotCycles(): Promise<Finding[]> {
|
||||
let entries: string[] = [];
|
||||
try { entries = await readdir(BOT_CYCLES_DIR); }
|
||||
catch { return []; }
|
||||
|
||||
const jsonFiles = entries.filter(e => e.endsWith(".json"));
|
||||
if (jsonFiles.length === 0) return [];
|
||||
|
||||
// Sort by mtime desc, take most recent N
|
||||
const withStat = await Promise.all(
|
||||
jsonFiles.map(async name => {
|
||||
try { return { name, mtime: (await stat(join(BOT_CYCLES_DIR, name))).mtimeMs }; }
|
||||
catch { return { name, mtime: 0 }; }
|
||||
}),
|
||||
);
|
||||
const recent = withStat.sort((a, b) => b.mtime - a.mtime).slice(0, MAX_BOT_CYCLE_FILES);
|
||||
|
||||
const outcomes: Record<string, number> = {};
|
||||
for (const { name } of recent) {
|
||||
try {
|
||||
const r = JSON.parse(await readFile(join(BOT_CYCLES_DIR, name), "utf8"));
|
||||
const o = String(r.outcome ?? "unknown");
|
||||
outcomes[o] = (outcomes[o] ?? 0) + 1;
|
||||
} catch { /* skip */ }
|
||||
}
|
||||
|
||||
const summary = Object.entries(outcomes)
|
||||
.sort((a, b) => b[1] - a[1])
|
||||
.map(([k, v]) => `${k}=${v}`)
|
||||
.join(", ");
|
||||
|
||||
const failCount = (outcomes["tests_failed"] ?? 0) + (outcomes["apply_failed"] ?? 0) + (outcomes["model_failed"] ?? 0);
|
||||
return [{
|
||||
check: "kb_query",
|
||||
severity: failCount > recent.length / 2 ? "warn" : "info",
|
||||
summary: `KB: bot recorded ${recent.length} recent cycles — ${summary || "no outcomes parsed"}`,
|
||||
evidence: [
|
||||
`dir: ${BOT_CYCLES_DIR}`,
|
||||
`fail-class total: ${failCount} / ${recent.length}`,
|
||||
],
|
||||
}];
|
||||
}
|
||||
|
||||
async function checkObserverStream(): Promise<Finding[]> {
|
||||
const ops = await tailJsonl<any>(OBSERVER_OPS, TAIL_LINES);
|
||||
if (ops.length === 0) return [];
|
||||
const failures = ops.filter(o => o.ok === false).length;
|
||||
return [{
|
||||
check: "kb_query",
|
||||
severity: "info",
|
||||
summary: `KB: observer stream ${ops.length} recent ops, ${failures} failures`,
|
||||
evidence: [
|
||||
`source: ${OBSERVER_OPS}`,
|
||||
`by source: ${observerBySource(ops)}`,
|
||||
],
|
||||
}];
|
||||
}
|
||||
|
||||
function observerBySource(ops: any[]): string {
|
||||
const c: Record<string, number> = {};
|
||||
for (const o of ops) {
|
||||
const s = String(o.source ?? "unknown");
|
||||
c[s] = (c[s] ?? 0) + 1;
|
||||
}
|
||||
return Object.entries(c).sort((a, b) => b[1] - a[1]).map(([k, v]) => `${k}=${v}`).join(", ") || "empty";
|
||||
}
|
||||
|
||||
// Audit-facts — reads data/_kb/audit_facts.jsonl (populated by every
|
||||
// curated inference run via llm_team's extract pipeline). Each row
|
||||
// has arrays: facts, entities, relationships. We explode entities and
|
||||
// aggregate them across PRs using kb_index. An entity seen in 3+ PRs
|
||||
// is a "core system entity" — we surface the top N as info context.
|
||||
//
|
||||
// Filters out short names (<3 chars, likely qwen2.5 truncation
|
||||
// artifacts) and generic types ("string", "number") that would
|
||||
// otherwise dominate the ranking.
|
||||
const ENTITY_NAME_MIN_LEN = 3;
|
||||
const GENERIC_ENTITY_NAMES = new Set([
|
||||
"string", "number", "boolean", "any", "void", "unknown", "never",
|
||||
"object", "array", "function", "const", "let", "var", "true", "false",
|
||||
"null", "undefined", "promise", "map", "set", "record",
|
||||
]);
|
||||
|
||||
async function checkAuditFacts(): Promise<Finding[]> {
|
||||
// Read raw rows — each row has multiple entities, so we can't just
|
||||
// use aggregate() directly (it's one-signature-per-row). Explode
|
||||
// entities into (row, entity) pairs, then aggregate by entity name.
|
||||
let raw: string;
|
||||
try { raw = await (await import("node:fs/promises")).readFile(AUDIT_FACTS_JSONL, "utf8"); }
|
||||
catch { return []; }
|
||||
const lines = raw.split("\n").filter(l => l.length > 0);
|
||||
if (lines.length === 0) return [];
|
||||
|
||||
interface EntityRow { entity_key: string; pr_number: number; type: string; name: string; description: string }
|
||||
const entityRows: EntityRow[] = [];
|
||||
for (const line of lines.slice(-TAIL_LINES * 2)) {
|
||||
let row: any;
|
||||
try { row = JSON.parse(line); } catch { continue; }
|
||||
const prNum = Number(row?.pr_number);
|
||||
if (!Number.isFinite(prNum)) continue;
|
||||
for (const e of Array.isArray(row?.entities) ? row.entities : []) {
|
||||
const name = String(e?.name ?? "").trim();
|
||||
if (name.length < ENTITY_NAME_MIN_LEN) continue;
|
||||
if (GENERIC_ENTITY_NAMES.has(name.toLowerCase())) continue;
|
||||
entityRows.push({
|
||||
entity_key: name.toLowerCase(),
|
||||
pr_number: prNum,
|
||||
type: String(e?.type ?? "?"),
|
||||
name,
|
||||
description: String(e?.description ?? "").slice(0, 160),
|
||||
});
|
||||
}
|
||||
}
|
||||
if (entityRows.length === 0) return [];
|
||||
|
||||
// Aggregate manually — one key per entity name, distinct_scopes by PR.
|
||||
type Agg = { count: number; scopes: Set<number>; types: Set<string>; last_name: string; last_desc: string };
|
||||
const byEntity = new Map<string, Agg>();
|
||||
for (const r of entityRows) {
|
||||
const a = byEntity.get(r.entity_key) ?? {
|
||||
count: 0, scopes: new Set<number>(), types: new Set<string>(), last_name: "", last_desc: "",
|
||||
};
|
||||
a.count += 1;
|
||||
a.scopes.add(r.pr_number);
|
||||
a.types.add(r.type);
|
||||
a.last_name = r.name;
|
||||
a.last_desc = r.description;
|
||||
byEntity.set(r.entity_key, a);
|
||||
}
|
||||
|
||||
// Rank: require 2+ distinct PRs (same-PR entity-repeats don't count
|
||||
// as "cross-cutting"). Take the top 5 to avoid flooding the verdict.
|
||||
const ranked = Array.from(byEntity.entries())
|
||||
.filter(([_, a]) => a.scopes.size >= 2)
|
||||
.sort((a, b) => b[1].scopes.size - a[1].scopes.size || b[1].count - a[1].count)
|
||||
.slice(0, 5);
|
||||
|
||||
if (ranked.length === 0) {
|
||||
// Useful to know the KB is being populated — emit a single
|
||||
// summary so operators see fact extraction is alive.
|
||||
return [{
|
||||
check: "kb_query",
|
||||
severity: "info",
|
||||
summary: `audit_facts KB has ${entityRows.length} entity-observations across ${new Set(entityRows.map(r => r.pr_number)).size} PRs (no cross-PR recurrences yet)`,
|
||||
evidence: [`source: ${AUDIT_FACTS_JSONL}`],
|
||||
}];
|
||||
}
|
||||
|
||||
return ranked.map(([_, a]) => ({
|
||||
check: "kb_query" as const,
|
||||
severity: "info" as const,
|
||||
summary: `core entity \`${a.last_name}\` recurs in ${a.scopes.size} PRs (types: ${Array.from(a.types).join(",")})`,
|
||||
evidence: [
|
||||
`count=${a.count} distinct_PRs=${a.scopes.size}`,
|
||||
`description: ${a.last_desc.slice(0, 200)}`,
|
||||
`PRs: ${Array.from(a.scopes).sort((x, y) => x - y).join(",")}`,
|
||||
],
|
||||
}));
|
||||
}
|
||||
|
||||
// Audit-lessons — reads data/_kb/audit_lessons.jsonl (populated by
|
||||
// every audit's appendAuditLessons). Uses the shared kb_index
|
||||
// aggregator: groups by `signature`, distinct-scopes keyed by PR
|
||||
// number, severity from ratingSeverity(agg) which applies the
|
||||
// confidence × count rating (see kb_index.ts). This is the same
|
||||
// aggregation any other KB reader uses — shared discipline, not
|
||||
// per-check custom logic.
|
||||
async function checkAuditLessons(): Promise<Finding[]> {
|
||||
const bySig = await aggregate<any>(AUDIT_LESSONS_JSONL, {
|
||||
keyFn: (r) => r?.signature,
|
||||
scopeFn: (r) => (r?.pr_number !== undefined ? `pr-${r.pr_number}` : undefined),
|
||||
checkFn: (r) => r?.check,
|
||||
tailLimit: TAIL_LINES * 4,
|
||||
});
|
||||
if (bySig.size === 0) return [];
|
||||
|
||||
const findings: Finding[] = [];
|
||||
for (const [sig, agg] of bySig) {
|
||||
// Silent on first-ever occurrence — not yet signal.
|
||||
if (agg.count < 2) continue;
|
||||
const sev = ratingSeverity(agg);
|
||||
findings.push({
|
||||
check: "kb_query",
|
||||
severity: sev,
|
||||
summary: `recurring audit pattern (${agg.distinct_scopes} distinct PRs, ${agg.count} flaggings, conf=${agg.confidence.toFixed(2)}): ${agg.representative_summary.slice(0, 160)}`,
|
||||
evidence: [
|
||||
`signature=${sig}`,
|
||||
`checks: ${agg.checks.join(",")}`,
|
||||
`scopes: ${agg.scopes.slice(-6).join(",")}`,
|
||||
formatAgg(agg),
|
||||
],
|
||||
});
|
||||
}
|
||||
return findings;
|
||||
}
|
||||
|
||||
// Scrum-master reviews — the scrum pipeline writes one row per
|
||||
// accepted per-file review. We match reviews whose `file` matches
|
||||
// any path in the PR's diff, then surface the *preview* + which
|
||||
// model the escalation ladder had to reach. If the scrum-master
|
||||
// needed the 123B specialist or larger to resolve a file, that's
|
||||
// a meaningful signal about the code's complexity — and it's
|
||||
// surfaced to the PR without the auditor having to re-run the
|
||||
// escalation ladder itself.
|
||||
async function checkScrumReviews(prFiles: string[]): Promise<Finding[]> {
|
||||
const rows = await tailJsonl<any>(SCRUM_REVIEWS_JSONL, TAIL_LINES);
|
||||
if (rows.length === 0) return [];
|
||||
|
||||
// Match by exact file OR filename suffix — PR files arrive as
|
||||
// `auditor/audit.ts`-style relative paths; scrum stores the same.
|
||||
const norm = (p: string) => p.replace(/^\/+/, "").replace(/^home\/profit\/lakehouse\//, "");
|
||||
const prSet = new Set(prFiles.map(norm));
|
||||
|
||||
// Keep only the most recent review per file (last-wins).
|
||||
const latestByFile = new Map<string, any>();
|
||||
for (const r of rows) {
|
||||
const f = norm(String(r.file ?? ""));
|
||||
if (!f) continue;
|
||||
if (!prSet.has(f)) continue;
|
||||
latestByFile.set(f, r);
|
||||
}
|
||||
if (latestByFile.size === 0) return [];
|
||||
|
||||
const findings: Finding[] = [];
|
||||
for (const [file, r] of latestByFile) {
|
||||
const model = String(r.accepted_model ?? "?");
|
||||
const attempt = r.accepted_on_attempt ?? "?";
|
||||
const treeSplit = !!r.tree_split_fired;
|
||||
// Heuristic: if the scrum-master had to escalate past attempt 3,
|
||||
// or had to tree-split, that's context the PR reviewer should see.
|
||||
// Severity: info for low-escalation, warn if escalated far up
|
||||
// the ladder (cloud specialist required).
|
||||
const heavyEscalation = Number(attempt) >= 4;
|
||||
const sev: "warn" | "info" = heavyEscalation ? "warn" : "info";
|
||||
findings.push({
|
||||
check: "kb_query",
|
||||
severity: sev,
|
||||
summary: `scrum-master review for \`${file}\` — accepted on attempt ${attempt} by \`${model}\`${treeSplit ? " (tree-split)" : ""}`,
|
||||
evidence: [
|
||||
`reviewed_at: ${r.reviewed_at ?? "?"}`,
|
||||
`preview: ${String(r.suggestions_preview ?? "").slice(0, 300).replace(/\n/g, " ")}`,
|
||||
],
|
||||
});
|
||||
}
|
||||
return findings;
|
||||
}
|
||||
@@ -1,184 +0,0 @@
|
||||
// Static diff check — grep-style, no AST, no LLM. Looks for patterns
|
||||
// that are high-signal evidence of placeholder code.
|
||||
//
|
||||
// Findings are severity-graded:
|
||||
// block — explicit non-impl markers (unimplemented!, todo!,
|
||||
// panic!("not implemented"), throw new Error("not implemented"))
|
||||
// warn — TODO / FIXME / XXX / HACK comments on added lines,
|
||||
// new struct fields with no read-site anywhere in the diff,
|
||||
// suspiciously-empty function bodies ({ Ok(()) } / {} when
|
||||
// the commit message claims the fn "implements" something)
|
||||
// info — hardcoded "test" / "dummy" / "placeholder" strings in
|
||||
// added lines (could be real, just flag for inspection)
|
||||
//
|
||||
// Consumes: raw unified diff text from Gitea.
|
||||
|
||||
import type { Finding } from "../types.ts";
|
||||
|
||||
// Rust + TypeScript patterns that almost always indicate "this is
|
||||
// not actually implemented yet."
|
||||
const BLOCK_PATTERNS: Array<{ re: RegExp; why: string }> = [
|
||||
{ re: /\bunimplemented!\s*\(/, why: "unimplemented!() macro call" },
|
||||
{ re: /\btodo!\s*\(/, why: "todo!() macro call" },
|
||||
{ re: /panic!\s*\(\s*"(?:not implemented|TODO|not yet|unimpl)/i, why: "panic! with not-implemented message" },
|
||||
{ re: /throw\s+new\s+Error\s*\(\s*['"](?:not implemented|TODO|unimpl)/i, why: "throw Error 'not implemented'" },
|
||||
];
|
||||
|
||||
const WARN_COMMENT_PATTERNS: Array<{ re: RegExp; why: string }> = [
|
||||
{ re: /^\+.*\/\/\s*(TODO|FIXME|XXX|HACK)\b/i, why: "TODO/FIXME/XXX/HACK comment added" },
|
||||
{ re: /^\+.*#\s*(TODO|FIXME|XXX|HACK)\b/i, why: "TODO/FIXME/XXX/HACK comment added" },
|
||||
];
|
||||
|
||||
const INFO_HARDCODED_PATTERNS: Array<{ re: RegExp; why: string }> = [
|
||||
{ re: /"(?:placeholder|dummy|foobar|xxx|replaceme|changeme)"/i, why: "suspicious hardcoded string" },
|
||||
];
|
||||
|
||||
export function runStaticCheck(diff: string): Finding[] {
|
||||
const findings: Finding[] = [];
|
||||
|
||||
// Per-file walk: only look at ADDED lines (prefix '+' but not '+++'
|
||||
// which is the diff header).
|
||||
const perFile = splitDiffByFile(diff);
|
||||
|
||||
for (const [path, lines] of perFile) {
|
||||
// Skip diff bookkeeping + pure-delete files
|
||||
if (!lines.some(l => l.startsWith("+") && !l.startsWith("+++"))) continue;
|
||||
|
||||
// The auditor's own check files literally contain the BLOCK
|
||||
// patterns as regex definitions (BLOCK_PATTERNS in this file,
|
||||
// prompt examples in inference.ts). Skipping BLOCK scan on these
|
||||
// specific paths prevents the checker from self-flagging its own
|
||||
// string literals. WARN/INFO patterns still run — those genuinely
|
||||
// could indicate problems in the checker's own code (TODO
|
||||
// comments don't self-define).
|
||||
const isAuditorCheckerFile = path.startsWith("auditor/checks/") ||
|
||||
path.startsWith("auditor/fixtures/");
|
||||
|
||||
for (let idx = 0; idx < lines.length; idx++) {
|
||||
const line = lines[idx];
|
||||
if (!line.startsWith("+") || line.startsWith("+++")) continue;
|
||||
const added = line.slice(1);
|
||||
|
||||
if (!isAuditorCheckerFile) {
|
||||
for (const { re, why } of BLOCK_PATTERNS) {
|
||||
const m = added.match(re);
|
||||
if (m && typeof m.index === "number") {
|
||||
// Skip if the match sits inside a quoted string literal —
|
||||
// this is how rubric files (tests/real-world/*, prompt
|
||||
// templates) legitimately reference the patterns they
|
||||
// guard against, without actually executing them.
|
||||
if (isInsideQuotedString(added, m.index)) continue;
|
||||
findings.push({
|
||||
check: "static",
|
||||
severity: "block",
|
||||
summary: `${why} in ${path}`,
|
||||
evidence: [`${path}:+${idx + 1}: ${added.trim().slice(0, 160)}`],
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
for (const { re, why } of WARN_COMMENT_PATTERNS) {
|
||||
if (re.test(line)) {
|
||||
findings.push({
|
||||
check: "static",
|
||||
severity: "warn",
|
||||
summary: `${why} in ${path}`,
|
||||
evidence: [`${path}:+${idx + 1}: ${added.trim().slice(0, 160)}`],
|
||||
});
|
||||
}
|
||||
}
|
||||
for (const { re, why } of INFO_HARDCODED_PATTERNS) {
|
||||
if (re.test(added)) {
|
||||
findings.push({
|
||||
check: "static",
|
||||
severity: "info",
|
||||
summary: `${why} in ${path}`,
|
||||
evidence: [`${path}:+${idx + 1}: ${added.trim().slice(0, 160)}`],
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// "Field added but never read" heuristic — catches exactly the
|
||||
// Phase 45 DocRef placeholder pattern. Limited to the diff itself:
|
||||
// we're not doing a full-codebase grep here (too noisy; callers
|
||||
// elsewhere might exist). The point is: if NEITHER this diff nor
|
||||
// any other line in the diff reads the field, the PR is shipping
|
||||
// state without a consumer.
|
||||
const addedLines = lines.filter(l => l.startsWith("+") && !l.startsWith("+++"))
|
||||
.map(l => l.slice(1));
|
||||
const newFields = extractNewFields(addedLines);
|
||||
for (const field of newFields) {
|
||||
const readPattern = new RegExp(`[\\.:]\\s*${escape(field)}\\b|\\b${escape(field)}\\s*:`);
|
||||
// The definition line itself matches readPattern — filter it out
|
||||
// by requiring at least TWO lines in the diff mention the field
|
||||
// (one defines, one reads).
|
||||
const hits = addedLines.filter(l => readPattern.test(l));
|
||||
if (hits.length < 2) {
|
||||
findings.push({
|
||||
check: "static",
|
||||
severity: "warn",
|
||||
summary: `field '${field}' added in ${path} but no read-site in the diff — could be placeholder state without a consumer`,
|
||||
evidence: [`${path}: added '${field}' with no reader; rest of diff has ${hits.length - 1} mentions`],
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return findings;
|
||||
}
|
||||
|
||||
function splitDiffByFile(diff: string): Map<string, string[]> {
|
||||
const out = new Map<string, string[]>();
|
||||
let current: string | null = null;
|
||||
let buf: string[] = [];
|
||||
for (const line of diff.split(/\r?\n/)) {
|
||||
const m = line.match(/^diff --git a\/(\S+) b\/(\S+)/);
|
||||
if (m) {
|
||||
if (current) out.set(current, buf);
|
||||
current = m[2];
|
||||
buf = [];
|
||||
continue;
|
||||
}
|
||||
buf.push(line);
|
||||
}
|
||||
if (current) out.set(current, buf);
|
||||
return out;
|
||||
}
|
||||
|
||||
// Extract new `pub name: Type,` fields from added lines. Rust syntax.
|
||||
// Narrowly-scoped: only matches at the start of a trimmed line,
|
||||
// requires `pub ` prefix, ignores `pub fn` / `pub struct` / etc.
|
||||
function extractNewFields(addedLines: string[]): string[] {
|
||||
const fields = new Set<string>();
|
||||
for (const line of addedLines) {
|
||||
const t = line.trim();
|
||||
// pub NAME: Type,
|
||||
const m = t.match(/^pub\s+(?!fn\b|struct\b|enum\b|mod\b|use\b|trait\b|impl\b|const\b|static\b|type\b)(\w+)\s*:/);
|
||||
if (m) fields.add(m[1]);
|
||||
}
|
||||
return Array.from(fields);
|
||||
}
|
||||
|
||||
// True if `pos` falls inside a double- or single-quoted string on this
|
||||
// line (backtick template literals too). Walks left→right toggling the
|
||||
// "in quote" state on each unescaped quote. Good enough for single-
|
||||
// line matches; multi-line strings aren't parsed (they're extremely
|
||||
// rare in the patterns we're blocking on, and would require a proper
|
||||
// tokenizer to handle correctly).
|
||||
function isInsideQuotedString(line: string, pos: number): boolean {
|
||||
let inDouble = false, inSingle = false, inBacktick = false;
|
||||
for (let i = 0; i < pos; i++) {
|
||||
const c = line[i];
|
||||
const esc = i > 0 && line[i - 1] === "\\";
|
||||
if (esc) continue;
|
||||
if (c === '"' && !inSingle && !inBacktick) inDouble = !inDouble;
|
||||
else if (c === "'" && !inDouble && !inBacktick) inSingle = !inSingle;
|
||||
else if (c === "`" && !inDouble && !inSingle) inBacktick = !inBacktick;
|
||||
}
|
||||
return inDouble || inSingle || inBacktick;
|
||||
}
|
||||
|
||||
function escape(s: string): string {
|
||||
return s.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
|
||||
}
|
||||
@@ -1,210 +0,0 @@
|
||||
// Claim parser — reads commit messages + PR body, extracts ship-claims.
|
||||
//
|
||||
// A "ship-claim" is any phrase that asserts functionality is working,
|
||||
// tested, complete, or landed. These are the assertions the downstream
|
||||
// checks (static/dynamic/inference/kb) try to falsify.
|
||||
//
|
||||
// Heuristic approach (regex + strength grading) — intentionally NOT
|
||||
// using an LLM here. Reason: the inference check already asks a cloud
|
||||
// model "does this match the claim?". The parser's job is to surface
|
||||
// the claim substrates, not judge them. Over-engineering the parser
|
||||
// risks false-negatives when the cloud model was going to catch it
|
||||
// anyway.
|
||||
|
||||
import type { Claim, PrSnapshot } from "./types.ts";
|
||||
|
||||
// Strong claims: explicit end-to-end + verification vocabulary
|
||||
const STRONG_PATTERNS: RegExp[] = [
|
||||
/\bverified\s+(end[- ]to[- ]end|live|in\s+production|against)\b/i,
|
||||
/\btested\s+(live|end[- ]to[- ]end|against|with)\b/i,
|
||||
/\bworks\s+(end[- ]to[- ]end|live|in\s+production)\b/i,
|
||||
/\bproduction[- ]ready\b/i,
|
||||
/\bfully\s+(functional|wired|working)\b/i,
|
||||
/\bphase\s+\d+(\.\d+)?\s+(shipped|complete|done|landed)\b/i,
|
||||
/\bground\s+truth\b/i,
|
||||
/\bproven\b/i,
|
||||
];
|
||||
|
||||
// Moderate claims: asserted completion or pass but without the strong
|
||||
// verification qualifier.
|
||||
const MODERATE_PATTERNS: RegExp[] = [
|
||||
/\bshipped\b/i,
|
||||
/\blanded\b/i,
|
||||
/\bgreen\b/i,
|
||||
/\b(tests?\s+)?pass(ing|ed)\b/i,
|
||||
/\bcomplet(e|ed)\b/i,
|
||||
/\bdone\b/i,
|
||||
/\bwired\b/i,
|
||||
/\bfixed\b/i,
|
||||
/\bworks\b/i,
|
||||
];
|
||||
|
||||
// Weak claims: aspirational or hedged. Usually low-risk but recorded
|
||||
// for completeness.
|
||||
const WEAK_PATTERNS: RegExp[] = [
|
||||
/\bshould\s+work\b/i,
|
||||
/\bexpected\s+to\b/i,
|
||||
/\bintended\s+to\b/i,
|
||||
/\bwill\s+(work|handle|support)\b/i,
|
||||
/\bprobably\b/i,
|
||||
];
|
||||
|
||||
// Empirical claims: runtime measurements / observed outcomes that can't
|
||||
// be verified from a diff (only from the actual run that produced
|
||||
// them). Classifying as empirical lets the inference check skip
|
||||
// diff-verification and saves the ladder for falsifiable claims.
|
||||
//
|
||||
// Two classes share this bucket because they share the skip discipline:
|
||||
//
|
||||
// 1. Runtime metrics — "58 cloud calls", "306s end-to-end"
|
||||
// 2. History/proof refs — "verified on PR #8", "was flipping across runs"
|
||||
//
|
||||
// Both are assertions about state outside the current diff. The cloud
|
||||
// would flag them as "not backed" — but that's a false positive: the
|
||||
// proof lives in the referenced run, prior commit, or test output, not
|
||||
// in the added lines the cloud is reading.
|
||||
const EMPIRICAL_PATTERNS: RegExp[] = [
|
||||
// ─── Runtime metrics ───
|
||||
// Iteration / attempt counts: "6/6 iterations", "attempt 5", "accepted on attempt 3"
|
||||
/\b\d+\s*\/\s*\d+\s+(iterations?|attempts?|cycles?|runs?|shards?)\b/i,
|
||||
/\b(accepted|resolved|converged)\s+on\s+attempt\s+\d+\b/i,
|
||||
// Runtime metrics: "58 cloud calls", "306s end-to-end", "245s total", "5931 chars"
|
||||
/\b\d+\s+(cloud\s+)?calls?\b/i,
|
||||
/\b\d+\s*(ms|s|seconds?|minutes?|m)\s+(end[- ]to[- ]end|total|elapsed|duration)\b/i,
|
||||
/\b\d+\s+chars?\b.*\b(accepted|generated|produced)\b/i,
|
||||
// "escalated through N tiers", "N distinct models"
|
||||
/\bescalated\s+through\s+\d+\b/i,
|
||||
/\b\d+\s+distinct\s+(model|tier)s?\b/i,
|
||||
|
||||
// ─── History / proof references ───
|
||||
// "verified on PR #8", "verified end-to-end on PR 8", "tested against PR #4"
|
||||
// Require PR#N / commit-hash / "prior <word>" to avoid matching
|
||||
// "verified ... in production" (PR without \b-ish anchor previously
|
||||
// consumed "pr" of "production").
|
||||
/\bverified\s+(?:end[- ]to[- ]end\s+)?(?:on|against|in)\s+(?:PR\s*#?\d+|commit\s+[0-9a-f]{6,}|prior\s+\w+|the\s+\w+\s+audit)\b/i,
|
||||
/\btested\s+(?:against|in|on)\s+(?:PR\s*#?\d+|commit\s+[0-9a-f]{6,}|prior\s+\w+)\b/i,
|
||||
// Direct PR/commit references: "PR #8", "on PR 9", "from commit abc123"
|
||||
/\b(?:on|from|in|via|per)\s+PR\s*#?\d+\b/i,
|
||||
/\b(?:from|in|per|against)\s+commit\s+[0-9a-f]{6,}/i,
|
||||
// Observational descriptions of prior behavior: "was flipping", "was X before", "previously observed"
|
||||
/\b(?:was|were)\s+(?:flipping|drifting|inconsistent|non[- ]deterministic|creeping)\b/i,
|
||||
/\bpreviously\s+(?:observed|flagged|reported|seen|landed)\b/i,
|
||||
/\bused\s+to\s+(?:flip|fail|flag|reject|block)\b/i,
|
||||
/\bobserved\s+(?:in|during|on|across)\s+(?:PR|prior|\d+\s+(?:runs?|audits?))/i,
|
||||
// "flipping/drifting across N runs" — historical variance description
|
||||
/\b(?:flipping|drifting|varying|oscillating)\s+across\s+(?:\d+\s+)?(?:runs?|audits?|iterations?)\b/i,
|
||||
// "the proven X" referring to prior work (proven is a STRONG pattern
|
||||
// but in context "the proven FOO" is usually a historical reference,
|
||||
// not a fresh claim). We catch it here so the empirical skip wins.
|
||||
/\bthe\s+proven\s+(?:escalation\s+ladder|pipeline|flow|loop|tier|path)/i,
|
||||
// "from the 9-run test", "across the 5-run validation"
|
||||
/\b(?:from|across|in|during)\s+the\s+\d+[- ]run\s+(?:test|validation|probe|experiment)/i,
|
||||
];
|
||||
|
||||
export interface ParsedClaims {
|
||||
claims: Claim[];
|
||||
commits_scanned: number;
|
||||
}
|
||||
|
||||
export function parseClaims(pr: PrSnapshot): ParsedClaims {
|
||||
const claims: Claim[] = [];
|
||||
|
||||
// PR body — every matching line becomes a claim at location "pr_body:N"
|
||||
if (pr.body) {
|
||||
scanText(pr.body, "pr_body", pr.head_sha, claims);
|
||||
}
|
||||
|
||||
// Each commit message gets its own scan.
|
||||
for (const c of pr.commits) {
|
||||
if (!c.message) continue;
|
||||
scanText(c.message, `commit:${c.sha.slice(0, 8)}`, c.sha, claims);
|
||||
}
|
||||
|
||||
return { claims, commits_scanned: pr.commits.length };
|
||||
}
|
||||
|
||||
function scanText(text: string, location_prefix: string, commit_sha: string, out: Claim[]): void {
|
||||
const lines = text.split(/\r?\n/);
|
||||
for (let i = 0; i < lines.length; i++) {
|
||||
const line = lines[i];
|
||||
if (line.length < 3) continue;
|
||||
|
||||
// Empirical match wins over everything else — if a line ALSO
|
||||
// contains a moderate word like "complete", we still want to
|
||||
// classify it as empirical so the inference check doesn't ask
|
||||
// the cloud to prove "58 cloud calls" from the diff. Order:
|
||||
// empirical → strong → moderate → weak.
|
||||
const empirical = firstUnquotedMatch(line, EMPIRICAL_PATTERNS);
|
||||
if (empirical) {
|
||||
out.push({
|
||||
text: line.trim().slice(0, 200),
|
||||
commit_sha,
|
||||
location: `${location_prefix}:${i + 1}`,
|
||||
strength: "empirical",
|
||||
});
|
||||
continue;
|
||||
}
|
||||
const strong = firstUnquotedMatch(line, STRONG_PATTERNS);
|
||||
if (strong) {
|
||||
out.push({
|
||||
text: line.trim().slice(0, 200),
|
||||
commit_sha,
|
||||
location: `${location_prefix}:${i + 1}`,
|
||||
strength: "strong",
|
||||
});
|
||||
continue;
|
||||
}
|
||||
const moderate = firstUnquotedMatch(line, MODERATE_PATTERNS);
|
||||
if (moderate) {
|
||||
out.push({
|
||||
text: line.trim().slice(0, 200),
|
||||
commit_sha,
|
||||
location: `${location_prefix}:${i + 1}`,
|
||||
strength: "moderate",
|
||||
});
|
||||
continue;
|
||||
}
|
||||
const weak = firstUnquotedMatch(line, WEAK_PATTERNS);
|
||||
if (weak) {
|
||||
out.push({
|
||||
text: line.trim().slice(0, 200),
|
||||
commit_sha,
|
||||
location: `${location_prefix}:${i + 1}`,
|
||||
strength: "weak",
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Match a pattern only when its match position is NOT inside a quoted
|
||||
// string on the line. Mirrors the same guard in auditor/checks/static.ts
|
||||
// — the two files have the same false-positive class: PR authors
|
||||
// quote pattern examples in commit message bodies (e.g. `"Phase 45
|
||||
// shipped"` as a test example) and without this guard those quoted
|
||||
// references get flagged as fresh ship-claims. Only skips when the
|
||||
// match itself falls inside quotes; real (unquoted) uses of the same
|
||||
// vocabulary still classify correctly.
|
||||
function firstUnquotedMatch(text: string, patterns: RegExp[]): RegExp | null {
|
||||
for (const p of patterns) {
|
||||
const m = text.match(p);
|
||||
if (!m || typeof m.index !== "number") continue;
|
||||
if (isInsideQuotedString(text, m.index)) continue;
|
||||
return p;
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
// Walks left→right toggling in-quote state on each unescaped quote.
|
||||
// Good enough for single-line claims; multi-line strings aren't parsed.
|
||||
function isInsideQuotedString(line: string, pos: number): boolean {
|
||||
let inDouble = false, inSingle = false, inBacktick = false;
|
||||
for (let i = 0; i < pos; i++) {
|
||||
const c = line[i];
|
||||
const esc = i > 0 && line[i - 1] === "\\";
|
||||
if (esc) continue;
|
||||
if (c === '"' && !inSingle && !inBacktick) inDouble = !inDouble;
|
||||
else if (c === "'" && !inDouble && !inBacktick) inSingle = !inSingle;
|
||||
else if (c === "`" && !inDouble && !inSingle) inBacktick = !inBacktick;
|
||||
}
|
||||
return inDouble || inSingle || inBacktick;
|
||||
}
|
||||
@@ -1,271 +0,0 @@
|
||||
// fact_extractor — routes curated TEXT through llm_team_ui's
|
||||
// "knowledge extract facts" mode (mode=extract at /api/run).
|
||||
//
|
||||
// What it gives us: structured {facts, entities, relationships} from
|
||||
// whatever curated blob we send. Auditor sends the tree-split
|
||||
// inference scratchpad (the best distillation of what a PR changed).
|
||||
// Scrum_master will later send its accepted review bodies.
|
||||
//
|
||||
// Why route through llm_team and not just extract directly from our
|
||||
// own checks: llm_team's extract uses a local EXTRACTOR model
|
||||
// (qwen2.5) + a separate VERIFIER (gemma2). This cross-check is the
|
||||
// discipline J wants for knowledge going into the playbook — facts
|
||||
// go in only after a second model has rated them CORRECT /
|
||||
// UNVERIFIABLE. Fast (local models, ~10-20s), free, and matches the
|
||||
// codereview pattern J already trusts.
|
||||
//
|
||||
// SSE parsing: llm_team streams SSE events. We're only interested in
|
||||
// the final "response" event with role="final" + the extraction
|
||||
// response (role="extraction N"). Parse the JSON from the extractor's
|
||||
// response text.
|
||||
|
||||
const LLM_TEAM = process.env.LH_LLM_TEAM_URL ?? "http://localhost:5000";
|
||||
const EXTRACTOR = process.env.LH_FACT_EXTRACTOR ?? "qwen2.5:latest";
|
||||
const VERIFIER = process.env.LH_FACT_VERIFIER ?? "gemma2:latest";
|
||||
const EXTRACT_TIMEOUT_MS = 120_000;
|
||||
const PROJECT_CONTEXT_FILE = process.env.LH_AUDITOR_CONTEXT_FILE
|
||||
?? "/home/profit/lakehouse/docs/AUDITOR_CONTEXT.md";
|
||||
|
||||
let cachedContext: string | null = null;
|
||||
async function loadProjectContext(): Promise<string> {
|
||||
if (cachedContext !== null) return cachedContext;
|
||||
try {
|
||||
const { readFile } = await import("node:fs/promises");
|
||||
const raw = await readFile(PROJECT_CONTEXT_FILE, "utf8");
|
||||
// Cap at 4KB — anything past that is more noise than signal for
|
||||
// the extractor/verifier's attention budget.
|
||||
cachedContext = raw.slice(0, 4000);
|
||||
} catch {
|
||||
cachedContext = ""; // context file missing → extractor runs without preamble
|
||||
}
|
||||
return cachedContext;
|
||||
}
|
||||
|
||||
export interface Entity {
|
||||
name: string;
|
||||
type: string;
|
||||
description?: string;
|
||||
}
|
||||
|
||||
export interface Relationship {
|
||||
from: string;
|
||||
to: string;
|
||||
type: string;
|
||||
}
|
||||
|
||||
export interface ExtractedFacts {
|
||||
facts: string[];
|
||||
entities: Entity[];
|
||||
relationships: Relationship[];
|
||||
verification: string;
|
||||
extractor_model: string;
|
||||
verifier_model: string;
|
||||
source_preview: string;
|
||||
// Populated when the extract run completed server-side (llm_team
|
||||
// persists to its own team_runs; this is for our own cross-ref).
|
||||
llm_team_run_id?: number;
|
||||
extracted_at: string;
|
||||
// Per-fact verdicts from the verifier pass (CORRECT/INCORRECT/
|
||||
// UNVERIFIABLE/UNCHECKED). Aligned 1:1 with the *raw* fact list
|
||||
// pre-drop so operators can see which verdicts mapped to dropped
|
||||
// facts if needed.
|
||||
verifier_verdicts?: Array<"CORRECT" | "INCORRECT" | "UNVERIFIABLE" | "UNCHECKED">;
|
||||
facts_dropped_by_verifier?: number;
|
||||
error?: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Run the llm_team extract pipeline on `source` text. Returns
|
||||
* structured {facts, entities, relationships}.
|
||||
*
|
||||
* Returns an object with `error` set if the pipeline failed — never
|
||||
* throws, because fact extraction is best-effort enrichment (the
|
||||
* primary audit must not break if llm_team is down).
|
||||
*/
|
||||
export async function extractFacts(source: string): Promise<ExtractedFacts> {
|
||||
const base: ExtractedFacts = {
|
||||
facts: [],
|
||||
entities: [],
|
||||
relationships: [],
|
||||
verification: "",
|
||||
extractor_model: EXTRACTOR,
|
||||
verifier_model: VERIFIER,
|
||||
source_preview: source.slice(0, 240),
|
||||
extracted_at: new Date().toISOString(),
|
||||
};
|
||||
|
||||
// Prepend project context to the source so the extractor + verifier
|
||||
// know what codebase/framework these facts belong to. Without this,
|
||||
// the verifier marks most domain-specific facts as UNVERIFIABLE ("I
|
||||
// don't know what Lakehouse is"). With it, the verifier can CORRECT-
|
||||
// stamp facts that align with the stated architecture.
|
||||
const context = await loadProjectContext();
|
||||
const prompt = context.length > 0
|
||||
? `=== PROJECT CONTEXT (for grounding facts; do NOT extract facts from this section) ===\n${context}\n\n=== CONTENT TO EXTRACT FACTS FROM ===\n${source}`
|
||||
: source;
|
||||
|
||||
let resp: Response;
|
||||
try {
|
||||
resp = await fetch(`${LLM_TEAM}/api/run`, {
|
||||
method: "POST",
|
||||
headers: { "content-type": "application/json" },
|
||||
body: JSON.stringify({
|
||||
mode: "extract",
|
||||
prompt,
|
||||
extractor: EXTRACTOR,
|
||||
verifier: VERIFIER,
|
||||
source: "prompt",
|
||||
skip_cache: true, // cache by prompt would dedup identical
|
||||
// scratchpads, but we want fresh extraction
|
||||
// for per-audit facts; cheap since local.
|
||||
}),
|
||||
signal: AbortSignal.timeout(EXTRACT_TIMEOUT_MS),
|
||||
});
|
||||
} catch (e) {
|
||||
return { ...base, error: `fetch failed: ${(e as Error).message}` };
|
||||
}
|
||||
|
||||
if (!resp.ok) {
|
||||
const body = await resp.text().catch(() => "");
|
||||
return { ...base, error: `llm_team /api/run ${resp.status}: ${body.slice(0, 200)}` };
|
||||
}
|
||||
|
||||
// Stream SSE lines; collect the one extraction response + the run_saved event
|
||||
// so we can capture the team-runs ID for cross-ref.
|
||||
const decoder = new TextDecoder();
|
||||
const reader = resp.body?.getReader();
|
||||
if (!reader) return { ...base, error: "no response body" };
|
||||
|
||||
let buffer = "";
|
||||
let extractionText = "";
|
||||
let verifierText = "";
|
||||
let runId: number | undefined = undefined;
|
||||
|
||||
try {
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) break;
|
||||
buffer += decoder.decode(value, { stream: true });
|
||||
let nl: number;
|
||||
while ((nl = buffer.indexOf("\n\n")) >= 0) {
|
||||
const chunk = buffer.slice(0, nl);
|
||||
buffer = buffer.slice(nl + 2);
|
||||
const dataLine = chunk.split("\n").find(l => l.startsWith("data: "));
|
||||
if (!dataLine) continue;
|
||||
try {
|
||||
const ev = JSON.parse(dataLine.slice(6));
|
||||
if (ev.type === "response") {
|
||||
const role = String(ev.role ?? "");
|
||||
if (role.startsWith("extraction")) extractionText = String(ev.text ?? "");
|
||||
else if (role === "verifier") verifierText = String(ev.text ?? "");
|
||||
} else if (ev.type === "run_saved") {
|
||||
const id = Number(ev.run_id);
|
||||
if (Number.isFinite(id)) runId = id;
|
||||
}
|
||||
} catch { /* skip malformed SSE */ }
|
||||
}
|
||||
}
|
||||
} catch (e) {
|
||||
return { ...base, error: `SSE read failed: ${(e as Error).message}` };
|
||||
}
|
||||
|
||||
// Pull the JSON object out of extractionText (may be wrapped in ```json fences).
|
||||
const parsed = extractFirstJsonObject(extractionText);
|
||||
if (!parsed) {
|
||||
return { ...base, error: "extractor returned no parseable JSON", verification: verifierText };
|
||||
}
|
||||
|
||||
const rawFacts: string[] = Array.isArray(parsed.facts)
|
||||
? parsed.facts.slice(0, 50).map(String)
|
||||
: [];
|
||||
|
||||
// Parse the verifier's free-form prose into per-fact verdicts, then
|
||||
// drop any fact the verifier explicitly marked INCORRECT. Leave
|
||||
// UNVERIFIABLE in place: many of our extractions are domain-specific
|
||||
// (Lakehouse internals) and the verifier has no prior-knowledge
|
||||
// anchor, so UNVERIFIABLE is the expected verdict for new signal,
|
||||
// not a quality fail. This is verifier-gated persistence: drop only
|
||||
// what's affirmatively wrong, not what's novel.
|
||||
const verdicts = parseVerifierVerdicts(verifierText, rawFacts.length);
|
||||
const incorrectIdx = new Set<number>();
|
||||
verdicts.forEach((v, i) => { if (v === "INCORRECT") incorrectIdx.add(i); });
|
||||
const kept = rawFacts.filter((_, i) => !incorrectIdx.has(i));
|
||||
|
||||
return {
|
||||
...base,
|
||||
facts: kept,
|
||||
entities: Array.isArray(parsed.entities)
|
||||
? parsed.entities.slice(0, 30).map((e: any) => ({
|
||||
name: String(e?.name ?? ""),
|
||||
type: String(e?.type ?? ""),
|
||||
description: typeof e?.description === "string" ? e.description.slice(0, 240) : undefined,
|
||||
})).filter(e => e.name.length > 0)
|
||||
: [],
|
||||
relationships: Array.isArray(parsed.relationships)
|
||||
? parsed.relationships.slice(0, 30).map((r: any) => ({
|
||||
from: String(r?.from ?? ""),
|
||||
to: String(r?.to ?? ""),
|
||||
type: String(r?.type ?? ""),
|
||||
})).filter(r => r.from.length > 0 && r.to.length > 0)
|
||||
: [],
|
||||
verification: verifierText.slice(0, 1500),
|
||||
facts_dropped_by_verifier: incorrectIdx.size,
|
||||
verifier_verdicts: verdicts,
|
||||
llm_team_run_id: runId,
|
||||
};
|
||||
}
|
||||
|
||||
// Parse verifier's free-form output into a per-fact verdict array.
|
||||
// Gemma2 uses several formats depending on prompt mood:
|
||||
// Format A: **1.** claim... * **Verdict:** CORRECT
|
||||
// Format B: **1.** claim... * **CORRECT** (no "Verdict:" label)
|
||||
// Format C: 1. claim... CORRECT
|
||||
// Strategy: split on fact numbers, then find the first
|
||||
// CORRECT|INCORRECT|UNVERIFIABLE token in each section. Handles all
|
||||
// three formats without regex gymnastics.
|
||||
function parseVerifierVerdicts(
|
||||
verifierText: string,
|
||||
numFacts: number,
|
||||
): Array<"CORRECT" | "INCORRECT" | "UNVERIFIABLE" | "UNCHECKED"> {
|
||||
const out: Array<"CORRECT" | "INCORRECT" | "UNVERIFIABLE" | "UNCHECKED"> =
|
||||
Array(numFacts).fill("UNCHECKED");
|
||||
if (!verifierText) return out;
|
||||
|
||||
// Find each fact section start — "**N.**" or "N." at line start —
|
||||
// and slice out the content up to the NEXT fact number. Each section
|
||||
// gets scanned for the first CORRECT/INCORRECT/UNVERIFIABLE token.
|
||||
const starts: Array<{ idx: number; pos: number }> = [];
|
||||
const header = /(?:^|\n)\s*(?:\*\*)?(\d+)[.)]/g;
|
||||
for (const m of verifierText.matchAll(header)) {
|
||||
const factNum = Number(m[1]);
|
||||
if (!Number.isFinite(factNum)) continue;
|
||||
starts.push({ idx: factNum - 1, pos: m.index! });
|
||||
}
|
||||
for (let i = 0; i < starts.length; i++) {
|
||||
const s = starts[i];
|
||||
const end = i + 1 < starts.length ? starts[i + 1].pos : verifierText.length;
|
||||
if (s.idx < 0 || s.idx >= numFacts) continue;
|
||||
const section = verifierText.slice(s.pos, end);
|
||||
const v = section.match(/\b(CORRECT|INCORRECT|UNVERIFIABLE)\b/i);
|
||||
if (v) out[s.idx] = v[1].toUpperCase() as "CORRECT" | "INCORRECT" | "UNVERIFIABLE";
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
// Lift the first balanced JSON object out of (possibly fenced) text.
|
||||
// Same discipline as inference.ts::extractJson.
|
||||
function extractFirstJsonObject(text: string): any | null {
|
||||
const cleaned = text.replace(/^```(?:json)?\s*/im, "").replace(/```\s*$/im, "");
|
||||
let depth = 0, start = -1;
|
||||
for (let i = 0; i < cleaned.length; i++) {
|
||||
const c = cleaned[i];
|
||||
if (c === "{") { if (depth === 0) start = i; depth++; }
|
||||
else if (c === "}") {
|
||||
depth--;
|
||||
if (depth === 0 && start >= 0) {
|
||||
try { return JSON.parse(cleaned.slice(start, i + 1)); } catch { start = -1; }
|
||||
}
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
@@ -1,49 +0,0 @@
|
||||
// Standalone runner for auditor fixtures. Invoke:
|
||||
//
|
||||
// bun run auditor/fixtures/cli.ts hybrid_38_40_45
|
||||
//
|
||||
// Prints human-readable per-layer breakdown + JSON result so the
|
||||
// output can be captured by the dynamic check without re-running.
|
||||
|
||||
import { runHybridFixture, type FixtureResult } from "./hybrid_38_40_45.ts";
|
||||
|
||||
const fixtureArg = process.argv[2] ?? "hybrid_38_40_45";
|
||||
|
||||
let result: FixtureResult;
|
||||
switch (fixtureArg) {
|
||||
case "hybrid_38_40_45":
|
||||
result = await runHybridFixture();
|
||||
break;
|
||||
default:
|
||||
console.error(`unknown fixture: ${fixtureArg}`);
|
||||
process.exit(2);
|
||||
}
|
||||
|
||||
// Human-readable summary
|
||||
console.error(""); // blank line to stderr so stdout JSON stays clean
|
||||
console.error(`─── Fixture: ${result.fixture} ───`);
|
||||
console.error(` overall: ${result.overall.toUpperCase()}`);
|
||||
console.error(` shipped: [${result.shipped_phases.join(", ") || "—"}]`);
|
||||
console.error(` placeholder: [${result.placeholder_phases.join(", ") || "—"}]`);
|
||||
console.error("");
|
||||
for (const l of result.layers) {
|
||||
const mark = l.ok ? "✓" : "✗";
|
||||
const phaseStr = `Phase ${l.phase}`.padEnd(11);
|
||||
console.error(` ${mark} ${phaseStr} ${l.layer.padEnd(30)} ${String(l.latency_ms).padStart(6)}ms`);
|
||||
if (l.ok) {
|
||||
console.error(` ${l.evidence.slice(0, 200)}`);
|
||||
} else {
|
||||
console.error(` ERROR: ${l.error?.slice(0, 240) ?? "unknown"}`);
|
||||
}
|
||||
}
|
||||
console.error("");
|
||||
for (const n of result.notes) console.error(` note: ${n}`);
|
||||
console.error("");
|
||||
|
||||
// Machine-readable output on stdout.
|
||||
console.log(JSON.stringify(result, null, 2));
|
||||
|
||||
// Exit code reflects overall: 0 pass, 2 partial, 1 fail.
|
||||
// Dynamic check reads this AND the JSON; partial-pass is treated
|
||||
// as informative (some layers shipped), not blocking on its own.
|
||||
process.exit(result.overall === "pass" ? 0 : result.overall === "partial_pass" ? 2 : 1);
|
||||
@@ -1,335 +0,0 @@
|
||||
// The never-run hybrid test fixture. Exercises the full stack end-to-end
|
||||
// across Phase 38 (/v1/chat), Phase 40 (Langfuse tracing), Phase 45
|
||||
// slice 1 (DocRef on playbook), and Phase 45 slice 2 (context7 bridge
|
||||
// drift detection).
|
||||
//
|
||||
// Returns HONEST per-layer results. A layer that's unimplemented or
|
||||
// broken returns ok=false with a real error; overall verdict flags
|
||||
// which phases are genuinely shipped vs still placeholder.
|
||||
//
|
||||
// Deterministic: always runs the same payload so a run's output is
|
||||
// comparable across audits. Uses TEST-prefixed names so playbook
|
||||
// state doesn't get polluted with production-looking fixture data.
|
||||
//
|
||||
// Preconditions:
|
||||
// - gateway up on :3100
|
||||
// - Python sidecar up on :3200
|
||||
// - Ollama local + Ollama Cloud key loaded
|
||||
// - Langfuse up on :3001
|
||||
// - context7 bridge up on :3900 (manual-start per mcp-server/README)
|
||||
// If bridge isn't running, that layer reports ok=false with a
|
||||
// clear error — the fixture doesn't try to auto-start it.
|
||||
|
||||
import { readFile } from "node:fs/promises";
|
||||
|
||||
const GATEWAY = process.env.LH_GATEWAY_URL ?? "http://localhost:3100";
|
||||
const LANGFUSE = process.env.LANGFUSE_URL ?? "http://localhost:3001";
|
||||
const BRIDGE = process.env.CONTEXT7_BRIDGE_URL ?? "http://localhost:3900";
|
||||
const FIXTURE_ID = "auditor:hybrid_38_40_45:v1";
|
||||
const TEST_WORKER_NAME = "__auditor_test_worker__";
|
||||
const STALE_HASH = "stale-hash-for-drift-proof";
|
||||
// Unique-per-run identifiers so each fixture invocation hits the ADD
|
||||
// path in upsert_entry (not UPDATE/NOOP on a leftover from a prior
|
||||
// run). Catches state pollution + avoids interaction with task #12
|
||||
// (UPDATE branch silently drops doc_refs).
|
||||
const RUN_NONCE = Date.now().toString(36);
|
||||
const TEST_WORKER_NAME_VERSIONED = `__auditor_test_worker_${RUN_NONCE}__`;
|
||||
const TEST_OPERATION_VERSIONED = `TEST: fill: __auditor_${RUN_NONCE}__ x1 in TestCity, TC`;
|
||||
|
||||
export interface LayerResult {
|
||||
layer: string;
|
||||
phase: string;
|
||||
ok: boolean;
|
||||
latency_ms: number;
|
||||
evidence: string;
|
||||
error?: string;
|
||||
// Layer-specific numbers that downstream checks can aggregate.
|
||||
metrics?: Record<string, number | string | boolean>;
|
||||
}
|
||||
|
||||
export interface FixtureResult {
|
||||
fixture: string;
|
||||
ran_at: string;
|
||||
overall: "pass" | "partial_pass" | "fail";
|
||||
layers: LayerResult[];
|
||||
real_numbers: Record<string, number>;
|
||||
shipped_phases: string[];
|
||||
placeholder_phases: string[];
|
||||
notes: string[];
|
||||
}
|
||||
|
||||
// Wraps a layer's async body with timing + structured error capture.
|
||||
// On throw: returns ok=false with the error string as evidence, so
|
||||
// downstream verdict code treats it as a layer failure rather than
|
||||
// a fixture crash.
|
||||
async function measureLayer(
|
||||
layer: string,
|
||||
phase: string,
|
||||
body: () => Promise<{ evidence: string; metrics?: Record<string, number | string | boolean> }>,
|
||||
): Promise<LayerResult> {
|
||||
const t0 = Date.now();
|
||||
try {
|
||||
const { evidence, metrics } = await body();
|
||||
return {
|
||||
layer, phase, ok: true,
|
||||
latency_ms: Date.now() - t0,
|
||||
evidence, metrics,
|
||||
};
|
||||
} catch (e) {
|
||||
return {
|
||||
layer, phase, ok: false,
|
||||
latency_ms: Date.now() - t0,
|
||||
evidence: `FAILED: ${(e as Error).message.slice(0, 200)}`,
|
||||
error: (e as Error).message,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
async function langfuseAuth(): Promise<{ pk: string; sk: string } | null> {
|
||||
// Default credentials match mcp-server/tracing.ts. Same sources the
|
||||
// Rust side uses in gateway/src/v1/langfuse_trace.rs.
|
||||
let pk = process.env.LANGFUSE_PUBLIC_KEY ?? "pk-lf-staffing";
|
||||
let sk = process.env.LANGFUSE_SECRET_KEY ?? "sk-lf-staffing-secret";
|
||||
if (!pk || !sk) return null;
|
||||
return { pk, sk };
|
||||
}
|
||||
|
||||
export async function runHybridFixture(): Promise<FixtureResult> {
|
||||
const result: FixtureResult = {
|
||||
fixture: FIXTURE_ID,
|
||||
ran_at: new Date().toISOString(),
|
||||
overall: "fail",
|
||||
layers: [],
|
||||
real_numbers: {},
|
||||
shipped_phases: [],
|
||||
placeholder_phases: [],
|
||||
notes: [],
|
||||
};
|
||||
|
||||
// ========================================================================
|
||||
// Layer 1 — Phase 38: POST /v1/chat returns valid OpenAI shape
|
||||
// ========================================================================
|
||||
// Captured HERE, immediately before the chat layer runs, so layer 2's
|
||||
// Langfuse-trace filter uses the actual moment the chat call was
|
||||
// attempted — not the fixture start time. Earlier draft had a
|
||||
// meaningless ternary returning result.ran_at on both branches; the
|
||||
// LLM-Team codereview (2026-04-22) caught this and flagged it as a
|
||||
// false-negative window on traces created between fixture-start and
|
||||
// chat-fetch.
|
||||
const chat_request_sent_ms = Date.now();
|
||||
|
||||
const l1 = await measureLayer("phase38_chat", "38", async () => {
|
||||
const r = await fetch(`${GATEWAY}/v1/chat`, {
|
||||
method: "POST",
|
||||
headers: { "content-type": "application/json" },
|
||||
body: JSON.stringify({
|
||||
model: "qwen3.5:latest",
|
||||
messages: [
|
||||
{ role: "system", content: "Respond in 8 words or fewer." },
|
||||
{ role: "user", content: "Name one Docker command." },
|
||||
],
|
||||
max_tokens: 60,
|
||||
temperature: 0.2,
|
||||
}),
|
||||
signal: AbortSignal.timeout(60000),
|
||||
});
|
||||
if (!r.ok) throw new Error(`gateway ${r.status}: ${await r.text()}`);
|
||||
const j: any = await r.json();
|
||||
const content = j.choices?.[0]?.message?.content;
|
||||
if (typeof content !== "string" || content.length === 0) {
|
||||
throw new Error(`empty content — think-false default may not be wired: ${JSON.stringify(j).slice(0, 200)}`);
|
||||
}
|
||||
return {
|
||||
evidence: `content="${content.slice(0, 60)}" tokens=${j.usage?.total_tokens}`,
|
||||
metrics: {
|
||||
prompt_tokens: j.usage?.prompt_tokens ?? 0,
|
||||
completion_tokens: j.usage?.completion_tokens ?? 0,
|
||||
total_tokens: j.usage?.total_tokens ?? 0,
|
||||
content_nonempty: true,
|
||||
},
|
||||
};
|
||||
});
|
||||
result.layers.push(l1);
|
||||
|
||||
// ========================================================================
|
||||
// Layer 2 — Phase 40: Langfuse trace lands within 3s
|
||||
// ========================================================================
|
||||
// Fire-and-forget tracing means we need a brief sleep before query.
|
||||
await new Promise(res => setTimeout(res, 2500));
|
||||
const l2 = await measureLayer("phase40_langfuse_trace", "40", async () => {
|
||||
const auth = await langfuseAuth();
|
||||
if (!auth) throw new Error("Langfuse credentials not in env");
|
||||
const r = await fetch(`${LANGFUSE}/api/public/traces?limit=5&name=v1.chat:ollama`, {
|
||||
headers: { Authorization: `Basic ${btoa(`${auth.pk}:${auth.sk}`)}` },
|
||||
signal: AbortSignal.timeout(10000),
|
||||
});
|
||||
if (!r.ok) throw new Error(`langfuse ${r.status}: ${await r.text()}`);
|
||||
const j: any = await r.json();
|
||||
const items = Array.isArray(j.data) ? j.data : [];
|
||||
// Filter on the chat-request timestamp captured above. A Langfuse
|
||||
// trace must be newer than the moment we fired /v1/chat to plausibly
|
||||
// belong to our request. Using fixture start time (result.ran_at)
|
||||
// was wrong and could false-negative on slow fixtures.
|
||||
const recent = items.filter((t: any) => Date.parse(t.timestamp) >= chat_request_sent_ms);
|
||||
if (recent.length === 0) {
|
||||
throw new Error(`no v1.chat:ollama trace since ${new Date(chat_request_sent_ms).toISOString()} (${items.length} older traces visible, Langfuse reachable — tracing is not firing)`);
|
||||
}
|
||||
const trace = recent[0];
|
||||
return {
|
||||
evidence: `trace ${trace.id.slice(0, 8)} name=${trace.name} age=${Date.now() - Date.parse(trace.timestamp)}ms`,
|
||||
metrics: {
|
||||
trace_age_ms: Date.now() - Date.parse(trace.timestamp),
|
||||
trace_latency_reported: Number(trace.latency ?? 0),
|
||||
recent_trace_count: recent.length,
|
||||
},
|
||||
};
|
||||
});
|
||||
result.layers.push(l2);
|
||||
|
||||
// ========================================================================
|
||||
// Layer 3 — Phase 45 slice 1: /seed accepts doc_refs and persists them
|
||||
// ========================================================================
|
||||
// Seed a playbook with doc_refs referencing Docker at a stale hash.
|
||||
// Ignore the cleanup concern for v1 — this fixture pollutes the
|
||||
// in-memory playbook state; operator can retire test entries.
|
||||
const l3 = await measureLayer("phase45_seed_with_doc_refs", "45.1", async () => {
|
||||
const r = await fetch(`${GATEWAY}/vectors/playbook_memory/seed`, {
|
||||
method: "POST",
|
||||
headers: { "content-type": "application/json" },
|
||||
body: JSON.stringify({
|
||||
operation: TEST_OPERATION_VERSIONED,
|
||||
approach: "auditor hybrid fixture run",
|
||||
context: "doc_refs integration test — retire after audit",
|
||||
endorsed_names: [TEST_WORKER_NAME_VERSIONED],
|
||||
append: true,
|
||||
doc_refs: [
|
||||
{
|
||||
tool: "docker",
|
||||
version_seen: "24.0.7",
|
||||
snippet_hash: STALE_HASH,
|
||||
source_url: "https://context7.com/docker/docs",
|
||||
seen_at: new Date().toISOString(),
|
||||
},
|
||||
],
|
||||
}),
|
||||
signal: AbortSignal.timeout(30000),
|
||||
});
|
||||
if (!r.ok) throw new Error(`seed ${r.status}: ${await r.text()}`);
|
||||
const j: any = await r.json();
|
||||
// Verify doc_refs actually persisted — read state.json directly.
|
||||
// That's the canonical source of truth; the seed response may or
|
||||
// may not echo doc_refs back.
|
||||
const state_raw = await readFile("/home/profit/lakehouse/data/_playbook_memory/state.json", "utf8");
|
||||
const state = JSON.parse(state_raw);
|
||||
const entries = state.entries ?? [];
|
||||
const ours = entries.find((e: any) => (e.endorsed_names ?? []).includes(TEST_WORKER_NAME_VERSIONED));
|
||||
if (!ours) throw new Error(`seeded entry not found in state.json (worker ${TEST_WORKER_NAME_VERSIONED} not present)`);
|
||||
const docRefs = ours.doc_refs ?? [];
|
||||
if (docRefs.length === 0) {
|
||||
throw new Error("entry saved but doc_refs field is empty — the field accepted by the API isn't being persisted");
|
||||
}
|
||||
return {
|
||||
evidence: `playbook ${ours.playbook_id.slice(0, 16)} persisted with doc_refs.length=${docRefs.length}; first tool=${docRefs[0]?.tool} hash=${docRefs[0]?.snippet_hash}`,
|
||||
metrics: {
|
||||
doc_refs_stored: docRefs.length,
|
||||
entries_after: j.entries_after ?? -1,
|
||||
playbook_id: ours.playbook_id,
|
||||
mode: j.outcome?.mode ?? "unknown",
|
||||
},
|
||||
};
|
||||
});
|
||||
result.layers.push(l3);
|
||||
|
||||
// ========================================================================
|
||||
// Layer 4 — Phase 45 slice 2: context7 bridge detects drift vs stale hash
|
||||
// ========================================================================
|
||||
const l4 = await measureLayer("phase45_bridge_diff", "45.2", async () => {
|
||||
// Health check first — if bridge isn't running, fail with a clear
|
||||
// message rather than mysterious connection refused.
|
||||
try {
|
||||
const h = await fetch(`${BRIDGE}/health`, { signal: AbortSignal.timeout(3000) });
|
||||
if (!h.ok) throw new Error(`bridge /health ${h.status}`);
|
||||
} catch (e) {
|
||||
throw new Error(`context7 bridge at ${BRIDGE} is not reachable (bridge is manual-start; run 'bun run mcp-server/context7_bridge.ts'): ${(e as Error).message}`);
|
||||
}
|
||||
const r = await fetch(`${BRIDGE}/docs/docker/diff?since=${encodeURIComponent(STALE_HASH)}`, {
|
||||
signal: AbortSignal.timeout(30000),
|
||||
});
|
||||
if (!r.ok) throw new Error(`bridge diff ${r.status}: ${await r.text()}`);
|
||||
const j: any = await r.json();
|
||||
if (j.drifted !== true) {
|
||||
throw new Error(`expected drifted=true against ${STALE_HASH}; got ${JSON.stringify(j)}`);
|
||||
}
|
||||
return {
|
||||
evidence: `bridge confirms drift: current=${j.current_snippet_hash} previous=${j.previous_snippet_hash} upstream_updated=${j.last_updated_upstream}`,
|
||||
metrics: {
|
||||
drifted: true,
|
||||
current_hash_length: String(j.current_snippet_hash ?? "").length,
|
||||
library_id: j.library_id,
|
||||
},
|
||||
};
|
||||
});
|
||||
result.layers.push(l4);
|
||||
|
||||
// ========================================================================
|
||||
// Layer 5 — Phase 45 slice 3: /doc_drift/check/{id} endpoint
|
||||
// (EXPECTED TO FAIL — endpoint unimplemented)
|
||||
// ========================================================================
|
||||
const l5 = await measureLayer("phase45_slice3_drift_flag", "45.3", async () => {
|
||||
const pid = String(l3.metrics?.playbook_id ?? "");
|
||||
if (!pid) throw new Error("layer 3 did not produce a playbook_id to check");
|
||||
const r = await fetch(`${GATEWAY}/vectors/playbook_memory/doc_drift/check/${encodeURIComponent(pid)}`, {
|
||||
method: "POST",
|
||||
signal: AbortSignal.timeout(10000),
|
||||
});
|
||||
if (r.status === 404 || r.status === 405) {
|
||||
// This is the HONEST signal: the endpoint doesn't exist yet.
|
||||
// Fail this layer with a clear marker — not a silent pass.
|
||||
throw new Error(`endpoint unimplemented (${r.status}) — Phase 45 slice 3 is still placeholder. doc_refs are persisted (layer 3) and the bridge can answer drift questions (layer 4), but nothing ties them together into a playbook flag yet.`);
|
||||
}
|
||||
if (!r.ok) throw new Error(`drift/check ${r.status}: ${await r.text()}`);
|
||||
const j: any = await r.json();
|
||||
// If the endpoint DID exist, assert the flag got set.
|
||||
if (!j.flagged) {
|
||||
throw new Error(`endpoint responded but flag not set: ${JSON.stringify(j).slice(0, 200)}`);
|
||||
}
|
||||
return {
|
||||
evidence: `playbook flagged: ${JSON.stringify(j).slice(0, 160)}`,
|
||||
metrics: { flagged: true },
|
||||
};
|
||||
});
|
||||
result.layers.push(l5);
|
||||
|
||||
// ========================================================================
|
||||
// Verdict assembly
|
||||
// ========================================================================
|
||||
for (const l of result.layers) {
|
||||
if (l.ok) result.shipped_phases.push(l.phase);
|
||||
else result.placeholder_phases.push(l.phase);
|
||||
result.real_numbers[`${l.layer}_latency_ms`] = l.latency_ms;
|
||||
if (l.metrics) {
|
||||
for (const [k, v] of Object.entries(l.metrics)) {
|
||||
if (typeof v === "number") result.real_numbers[`${l.layer}.${k}`] = v;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const anyFail = result.placeholder_phases.length > 0;
|
||||
const allPass = result.layers.every(l => l.ok);
|
||||
result.overall = allPass ? "pass" : (result.shipped_phases.length > 0 ? "partial_pass" : "fail");
|
||||
|
||||
if (!l5.ok) {
|
||||
result.notes.push(
|
||||
"Phase 45 slice 3 (doc_drift/check endpoint) is the expected-fail layer. " +
|
||||
"Layer 5 failing with a 404/405 is the CORRECT honest signal — don't mask this " +
|
||||
"to get a green result. When slice 3 ships, layer 5 flips to ok=true automatically."
|
||||
);
|
||||
}
|
||||
if (anyFail) {
|
||||
result.notes.push(
|
||||
`${result.placeholder_phases.length} layer(s) failed. Check evidence per layer for specifics.`
|
||||
);
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
@@ -1,145 +0,0 @@
|
||||
// Gitea API client. Minimal surface — only what the auditor needs:
|
||||
// list open PRs, get commits + files for a PR, fetch a diff, post a
|
||||
// commit status, post a review.
|
||||
//
|
||||
// Auth: reads PAT from ~/.git-credentials (set up by the credential
|
||||
// helper flow in 2026-04-22 session). Gitea's "token" auth scheme
|
||||
// matches what `git fetch` is already using.
|
||||
|
||||
import { readFile } from "node:fs/promises";
|
||||
import type { PrSnapshot } from "./types.ts";
|
||||
|
||||
const HOST = process.env.GITEA_HOST ?? "https://git.agentview.dev";
|
||||
const OWNER = "profit";
|
||||
const REPO = "lakehouse";
|
||||
const CRED_FILE = "/home/profit/.git-credentials";
|
||||
|
||||
let cachedPat: string | null = null;
|
||||
|
||||
async function getPat(): Promise<string> {
|
||||
if (cachedPat) return cachedPat;
|
||||
const raw = await readFile(CRED_FILE, "utf8");
|
||||
for (const line of raw.split("\n")) {
|
||||
const m = line.match(/^https:\/\/[^:]+:([^@]+)@git\.agentview\.dev/);
|
||||
if (m) { cachedPat = m[1]; return m[1]; }
|
||||
}
|
||||
throw new Error(`no Gitea PAT in ${CRED_FILE}`);
|
||||
}
|
||||
|
||||
async function giteaFetch(path: string, init: RequestInit = {}): Promise<Response> {
|
||||
const pat = await getPat();
|
||||
const url = `${HOST}/api/v1${path}`;
|
||||
const headers = new Headers(init.headers);
|
||||
headers.set("Authorization", `token ${pat}`);
|
||||
if (init.body && !headers.has("content-type")) {
|
||||
headers.set("content-type", "application/json");
|
||||
}
|
||||
return fetch(url, { ...init, headers, signal: AbortSignal.timeout(20000) });
|
||||
}
|
||||
|
||||
export async function listOpenPrs(): Promise<PrSnapshot[]> {
|
||||
const r = await giteaFetch(`/repos/${OWNER}/${REPO}/pulls?state=open&page=1&limit=50`);
|
||||
if (!r.ok) throw new Error(`listOpenPrs ${r.status}: ${await r.text()}`);
|
||||
const rows = (await r.json()) as any[];
|
||||
return Promise.all(rows.map(row => snapshotFromPr(row)));
|
||||
}
|
||||
|
||||
export async function getPrSnapshot(num: number): Promise<PrSnapshot> {
|
||||
const r = await giteaFetch(`/repos/${OWNER}/${REPO}/pulls/${num}`);
|
||||
if (!r.ok) throw new Error(`getPr ${num} ${r.status}: ${await r.text()}`);
|
||||
return snapshotFromPr((await r.json()) as any);
|
||||
}
|
||||
|
||||
async function snapshotFromPr(row: any): Promise<PrSnapshot> {
|
||||
const num = row.number;
|
||||
const commitsResp = await giteaFetch(`/repos/${OWNER}/${REPO}/pulls/${num}/commits`);
|
||||
const commits = commitsResp.ok ? ((await commitsResp.json()) as any[]) : [];
|
||||
const filesResp = await giteaFetch(`/repos/${OWNER}/${REPO}/pulls/${num}/files`);
|
||||
const files = filesResp.ok ? ((await filesResp.json()) as any[]) : [];
|
||||
return {
|
||||
number: num,
|
||||
head_sha: row.head?.sha ?? "",
|
||||
base_sha: row.base?.sha ?? "",
|
||||
title: row.title ?? "",
|
||||
body: row.body ?? "",
|
||||
state: row.state === "open" ? "open" : (row.merged ? "merged" : "closed"),
|
||||
author: row.user?.login ?? "",
|
||||
commits: commits.map(c => ({
|
||||
sha: (c.sha ?? "").slice(0, 12),
|
||||
message: c.commit?.message ?? "",
|
||||
author: c.commit?.author?.name ?? "",
|
||||
})),
|
||||
files: files.map(f => ({
|
||||
path: f.filename ?? "",
|
||||
additions: f.additions ?? 0,
|
||||
deletions: f.deletions ?? 0,
|
||||
})),
|
||||
};
|
||||
}
|
||||
|
||||
/// Returns the unified diff text of the PR. Used by static checks.
|
||||
export async function getPrDiff(num: number): Promise<string> {
|
||||
const r = await giteaFetch(`/repos/${OWNER}/${REPO}/pulls/${num}.diff`);
|
||||
if (!r.ok) throw new Error(`getDiff ${num} ${r.status}: ${await r.text()}`);
|
||||
return await r.text();
|
||||
}
|
||||
|
||||
/// Hard-block mechanism: post a failing commit status on the PR head
|
||||
/// SHA. Branch protection (if enabled on `main`) treats this as a
|
||||
/// required-check fail and prevents merge. The description is shown
|
||||
/// in the Gitea UI next to the red X.
|
||||
export async function postCommitStatus(args: {
|
||||
sha: string;
|
||||
state: "success" | "pending" | "failure" | "error";
|
||||
context: string;
|
||||
description: string;
|
||||
target_url?: string;
|
||||
}): Promise<void> {
|
||||
const r = await giteaFetch(`/repos/${OWNER}/${REPO}/statuses/${args.sha}`, {
|
||||
method: "POST",
|
||||
body: JSON.stringify({
|
||||
state: args.state,
|
||||
context: args.context,
|
||||
description: args.description.slice(0, 140),
|
||||
target_url: args.target_url ?? "",
|
||||
}),
|
||||
});
|
||||
if (!r.ok) throw new Error(`postCommitStatus ${r.status}: ${await r.text()}`);
|
||||
}
|
||||
|
||||
/// Post a review comment. Gitea typically blocks self-review
|
||||
/// (author posting a review on their own PR). Prefer
|
||||
/// `postIssueComment` when running with the author's PAT.
|
||||
export async function postReview(args: {
|
||||
pr_number: number;
|
||||
commit_id: string;
|
||||
body: string;
|
||||
event: "APPROVE" | "REQUEST_CHANGES" | "COMMENT";
|
||||
}): Promise<void> {
|
||||
const r = await giteaFetch(`/repos/${OWNER}/${REPO}/pulls/${args.pr_number}/reviews`, {
|
||||
method: "POST",
|
||||
body: JSON.stringify({
|
||||
commit_id: args.commit_id,
|
||||
body: args.body,
|
||||
event: args.event,
|
||||
}),
|
||||
});
|
||||
if (!r.ok) throw new Error(`postReview ${r.status}: ${await r.text()}`);
|
||||
}
|
||||
|
||||
/// Plain issue comment. Works for the auditor's own PAT because
|
||||
/// Gitea allows authors to comment on their own PRs (just not
|
||||
/// review them). Auditor uses this for the reasoning body; the
|
||||
/// actual block signal is the commit status.
|
||||
export async function postIssueComment(args: {
|
||||
pr_number: number;
|
||||
body: string;
|
||||
}): Promise<{ id: number; html_url: string }> {
|
||||
const r = await giteaFetch(`/repos/${OWNER}/${REPO}/issues/${args.pr_number}/comments`, {
|
||||
method: "POST",
|
||||
body: JSON.stringify({ body: args.body }),
|
||||
});
|
||||
if (!r.ok) throw new Error(`postIssueComment ${r.status}: ${await r.text()}`);
|
||||
const j = await r.json() as any;
|
||||
return { id: j.id, html_url: j.html_url };
|
||||
}
|
||||
@@ -1,147 +0,0 @@
|
||||
// Auditor poller — the top-level entry. Polls Gitea for open PRs on
|
||||
// a fixed interval, dedupes by head SHA, runs audit + posts verdict
|
||||
// for each new (pr, sha) pair.
|
||||
//
|
||||
// Run manually:
|
||||
// bun run auditor/index.ts
|
||||
//
|
||||
// Stop:
|
||||
// touch auditor.paused (skips next cycle)
|
||||
// pkill -f auditor/index.ts (kills in-flight)
|
||||
//
|
||||
// State:
|
||||
// data/_auditor/state.json — last-audited SHA per PR
|
||||
// data/_auditor/verdicts/{id}.json — per-run verdict records
|
||||
//
|
||||
// This entry runs forever. A systemd unit would wrap it once the
|
||||
// workflow is trusted (same pattern as mcp-server, observer).
|
||||
|
||||
import { readFile, writeFile, mkdir, access } from "node:fs/promises";
|
||||
import { listOpenPrs } from "./gitea.ts";
|
||||
import { auditPr } from "./audit.ts";
|
||||
|
||||
const POLL_INTERVAL_MS = 90_000; // 90s — enough budget for audit runs to complete
|
||||
const PAUSE_FILE = "/home/profit/lakehouse/auditor.paused";
|
||||
const STATE_FILE = "/home/profit/lakehouse/data/_auditor/state.json";
|
||||
|
||||
interface State {
|
||||
// Map: PR number → last-audited head SHA. Lets us dedupe audits
|
||||
// across restarts (poller can crash/restart without re-auditing
|
||||
// all open PRs from scratch).
|
||||
last_audited: Record<string, string>;
|
||||
started_at: string;
|
||||
cycles_total: number;
|
||||
cycles_skipped_paused: number;
|
||||
audits_run: number;
|
||||
last_cycle_at?: string;
|
||||
}
|
||||
|
||||
async function fileExists(path: string): Promise<boolean> {
|
||||
try { await access(path); return true; } catch { return false; }
|
||||
}
|
||||
|
||||
async function loadState(): Promise<State> {
|
||||
try {
|
||||
const raw = await readFile(STATE_FILE, "utf8");
|
||||
const s = JSON.parse(raw);
|
||||
return {
|
||||
last_audited: s.last_audited ?? {},
|
||||
started_at: s.started_at ?? new Date().toISOString(),
|
||||
cycles_total: s.cycles_total ?? 0,
|
||||
cycles_skipped_paused: s.cycles_skipped_paused ?? 0,
|
||||
audits_run: s.audits_run ?? 0,
|
||||
last_cycle_at: s.last_cycle_at,
|
||||
};
|
||||
} catch {
|
||||
return {
|
||||
last_audited: {},
|
||||
started_at: new Date().toISOString(),
|
||||
cycles_total: 0,
|
||||
cycles_skipped_paused: 0,
|
||||
audits_run: 0,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
async function saveState(s: State): Promise<void> {
|
||||
await mkdir("/home/profit/lakehouse/data/_auditor", { recursive: true });
|
||||
await writeFile(STATE_FILE, JSON.stringify(s, null, 2));
|
||||
}
|
||||
|
||||
async function runCycle(state: State): Promise<State> {
|
||||
state.cycles_total += 1;
|
||||
state.last_cycle_at = new Date().toISOString();
|
||||
|
||||
if (await fileExists(PAUSE_FILE)) {
|
||||
state.cycles_skipped_paused += 1;
|
||||
console.log(`[auditor] cycle ${state.cycles_total}: paused (touch ${PAUSE_FILE} exists)`);
|
||||
return state;
|
||||
}
|
||||
|
||||
let prs;
|
||||
try {
|
||||
prs = await listOpenPrs();
|
||||
} catch (e) {
|
||||
console.error(`[auditor] listOpenPrs failed: ${(e as Error).message}`);
|
||||
return state;
|
||||
}
|
||||
|
||||
console.log(`[auditor] cycle ${state.cycles_total}: ${prs.length} open PR(s)`);
|
||||
|
||||
for (const pr of prs) {
|
||||
const last = state.last_audited[String(pr.number)];
|
||||
if (last === pr.head_sha) {
|
||||
console.log(`[auditor] skip PR #${pr.number} (SHA ${pr.head_sha.slice(0, 8)} already audited)`);
|
||||
continue;
|
||||
}
|
||||
console.log(`[auditor] audit PR #${pr.number} (${pr.head_sha.slice(0, 8)}) — ${pr.title.slice(0, 60)}`);
|
||||
try {
|
||||
// Skip dynamic by default: it mutates live playbook state and
|
||||
// re-runs on every PR update would pollute quickly. Operator
|
||||
// can run dynamic via `bun run auditor/fixtures/cli.ts` manually
|
||||
// OR set LH_AUDITOR_RUN_DYNAMIC=1 to opt in.
|
||||
const run_dynamic = process.env.LH_AUDITOR_RUN_DYNAMIC === "1";
|
||||
const verdict = await auditPr(pr, {
|
||||
skip_dynamic: !run_dynamic,
|
||||
skip_inference: process.env.LH_AUDITOR_SKIP_INFERENCE === "1",
|
||||
});
|
||||
console.log(`[auditor] verdict=${verdict.overall} findings=${verdict.metrics.findings_total} (block=${verdict.metrics.findings_block} warn=${verdict.metrics.findings_warn})`);
|
||||
state.last_audited[String(pr.number)] = pr.head_sha;
|
||||
state.audits_run += 1;
|
||||
} catch (e) {
|
||||
console.error(`[auditor] audit failed: ${(e as Error).message}`);
|
||||
}
|
||||
}
|
||||
|
||||
return state;
|
||||
}
|
||||
|
||||
async function main(): Promise<void> {
|
||||
console.log(`[auditor] starting poller — interval ${POLL_INTERVAL_MS / 1000}s`);
|
||||
console.log(`[auditor] pause file: ${PAUSE_FILE}`);
|
||||
console.log(`[auditor] state file: ${STATE_FILE}`);
|
||||
|
||||
let state = await loadState();
|
||||
console.log(`[auditor] loaded state: ${Object.keys(state.last_audited).length} PRs previously audited, ${state.cycles_total} cycles so far`);
|
||||
|
||||
// Single-shot mode for CLI testing: `bun run auditor/index.ts --once`
|
||||
const once = process.argv.includes("--once");
|
||||
if (once) {
|
||||
state = await runCycle(state);
|
||||
await saveState(state);
|
||||
console.log(`[auditor] single-shot complete. total audits: ${state.audits_run}`);
|
||||
return;
|
||||
}
|
||||
|
||||
// Loop.
|
||||
while (true) {
|
||||
state = await runCycle(state);
|
||||
await saveState(state);
|
||||
await new Promise(res => setTimeout(res, POLL_INTERVAL_MS));
|
||||
}
|
||||
}
|
||||
|
||||
main().catch(e => {
|
||||
console.error("[auditor] fatal:", e);
|
||||
process.exit(1);
|
||||
});
|
||||
@@ -1,161 +0,0 @@
|
||||
// kb_index — generic on-the-fly aggregation over append-only JSONL
|
||||
// scratchpads (audit_lessons, scrum_reviews, outcomes, observer ops).
|
||||
//
|
||||
// The mem0 insight: raw rows are CHEAP and tell the full story, but
|
||||
// downstream prompts need a DEFINITION, not a log. A definition is
|
||||
// the aggregate: "this signature has fired N times across M distinct
|
||||
// scopes, first_seen=X, last_seen=Y, confidence=M/N."
|
||||
//
|
||||
// This library is the single shared aggregator. Every KB writer keeps
|
||||
// appending raw rows; every KB reader uses aggregate() instead of
|
||||
// tailing the raw stream. No second file to sync, no ADD/UPDATE/NOOP
|
||||
// routing — the stats roll up from the raw rows every time.
|
||||
//
|
||||
// Why this works past hundreds of runs:
|
||||
// - aggregate() is bounded by distinct_signatures, not total_rows.
|
||||
// - confidence = distinct_scopes / count — low for same-scope noise,
|
||||
// high for cross-scope patterns. Downstream severity ramps on
|
||||
// confidence × count, not raw count, so one unfixed PR can't
|
||||
// inflate its own recurrence score (the classic mem0 failure).
|
||||
// - rotation (later) moves old raw to archive files; aggregate()
|
||||
// can still read both to compute lifetime counts when needed.
|
||||
|
||||
import { readFile } from "node:fs/promises";
|
||||
|
||||
export interface AggregateRow {
|
||||
signature: string;
|
||||
count: number;
|
||||
distinct_scopes: number;
|
||||
first_seen: string;
|
||||
last_seen: string;
|
||||
confidence: number; // distinct_scopes / count — capped at 1.0
|
||||
representative_summary: string; // most-recent summary for this signature
|
||||
scopes: string[]; // up to 20 most-recent scopes for debugging
|
||||
checks: string[]; // distinct `check` values (audit_lessons-specific)
|
||||
}
|
||||
|
||||
export interface AggregateOptions<T> {
|
||||
/** How to extract the dedup key from a row. */
|
||||
keyFn: (row: T) => string | undefined;
|
||||
/** How to extract the "scope" — distinct scopes count gives confidence. */
|
||||
scopeFn: (row: T) => string | undefined;
|
||||
/** How to extract the timestamp (defaults to row.audited_at / row.reviewed_at / row.timestamp). */
|
||||
timeFn?: (row: T) => string | undefined;
|
||||
/** How to extract a representative summary (defaults to row.summary). */
|
||||
summaryFn?: (row: T) => string | undefined;
|
||||
/** Max rows to read from the JSONL tail; 0 = read all. */
|
||||
tailLimit?: number;
|
||||
/** Include per-row check field (for multi-check aggregates). */
|
||||
checkFn?: (row: T) => string | undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
* Read a JSONL file and produce the aggregate map keyed by signature.
|
||||
* Safe on missing or malformed files — returns empty map.
|
||||
*/
|
||||
export async function aggregate<T = any>(
|
||||
jsonlPath: string,
|
||||
opts: AggregateOptions<T>,
|
||||
): Promise<Map<string, AggregateRow>> {
|
||||
const out = new Map<string, AggregateRow>();
|
||||
let raw: string;
|
||||
try { raw = await readFile(jsonlPath, "utf8"); } catch { return out; }
|
||||
const lines = raw.split("\n").filter(l => l.length > 0);
|
||||
const sliceFrom = opts.tailLimit && opts.tailLimit > 0 ? Math.max(0, lines.length - opts.tailLimit) : 0;
|
||||
|
||||
const timeFn = opts.timeFn ?? ((r: any) => r?.audited_at ?? r?.reviewed_at ?? r?.timestamp ?? r?.ran_at);
|
||||
const summaryFn = opts.summaryFn ?? ((r: any) => r?.summary ?? r?.representative_summary);
|
||||
|
||||
// Per-signature scope tracking — need counts by scope to compute
|
||||
// distinct_scopes without double-counting a scope that appears 50
|
||||
// times. Using a Set<scope> per signature.
|
||||
const scopeSets = new Map<string, Set<string>>();
|
||||
const checkSets = new Map<string, Set<string>>();
|
||||
|
||||
for (let i = sliceFrom; i < lines.length; i++) {
|
||||
let row: T;
|
||||
try { row = JSON.parse(lines[i]) as T; } catch { continue; }
|
||||
const sig = opts.keyFn(row);
|
||||
if (!sig) continue;
|
||||
|
||||
let agg = out.get(sig);
|
||||
if (!agg) {
|
||||
agg = {
|
||||
signature: sig,
|
||||
count: 0,
|
||||
distinct_scopes: 0,
|
||||
first_seen: "",
|
||||
last_seen: "",
|
||||
confidence: 0,
|
||||
representative_summary: "",
|
||||
scopes: [],
|
||||
checks: [],
|
||||
};
|
||||
out.set(sig, agg);
|
||||
scopeSets.set(sig, new Set<string>());
|
||||
checkSets.set(sig, new Set<string>());
|
||||
}
|
||||
|
||||
agg.count += 1;
|
||||
|
||||
const scope = opts.scopeFn(row);
|
||||
if (scope !== undefined && scope !== null && scope !== "") {
|
||||
scopeSets.get(sig)!.add(String(scope));
|
||||
// Keep scopes array ordered by recency (newest wins — shift
|
||||
// oldest when at cap).
|
||||
const arr = agg.scopes;
|
||||
const s = String(scope);
|
||||
const existing = arr.indexOf(s);
|
||||
if (existing >= 0) arr.splice(existing, 1);
|
||||
arr.push(s);
|
||||
if (arr.length > 20) arr.shift();
|
||||
}
|
||||
|
||||
if (opts.checkFn) {
|
||||
const c = opts.checkFn(row);
|
||||
if (c) checkSets.get(sig)!.add(String(c));
|
||||
}
|
||||
|
||||
const t = timeFn(row);
|
||||
if (t) {
|
||||
if (!agg.first_seen || t < agg.first_seen) agg.first_seen = t;
|
||||
if (!agg.last_seen || t > agg.last_seen) agg.last_seen = t;
|
||||
}
|
||||
|
||||
const s = summaryFn(row);
|
||||
if (s) agg.representative_summary = String(s);
|
||||
}
|
||||
|
||||
// Finalize derived fields.
|
||||
for (const [sig, agg] of out) {
|
||||
const scopes = scopeSets.get(sig) ?? new Set<string>();
|
||||
agg.distinct_scopes = scopes.size;
|
||||
agg.confidence = agg.count > 0 ? Math.min(1, agg.distinct_scopes / agg.count) : 0;
|
||||
const checks = checkSets.get(sig);
|
||||
if (checks) agg.checks = Array.from(checks).sort();
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
/**
|
||||
* Severity policy derived from aggregate stats. The rating lives here
|
||||
* (not in each check) so all KB readers ramp severity consistently.
|
||||
*
|
||||
* - confidence × count product is the real signal.
|
||||
* - Low confidence (< 0.3) = same-scope noise → info regardless of count.
|
||||
* - Mid confidence (0.3-0.6) = mixed signal → warn at count ≥ 3.
|
||||
* - High confidence (> 0.6) with count ≥ 5 = block-worthy cross-cutting pattern.
|
||||
*
|
||||
* Callers can override by reading agg directly; this is the default
|
||||
* policy that matches the "don't escalate one unfixed PR" discipline.
|
||||
*/
|
||||
export function ratingSeverity(agg: AggregateRow): "info" | "warn" | "block" {
|
||||
if (agg.confidence >= 0.6 && agg.count >= 5) return "block";
|
||||
if (agg.confidence >= 0.3 && agg.count >= 3) return "warn";
|
||||
return "info";
|
||||
}
|
||||
|
||||
/** Human-friendly one-line summary of an aggregate row for finding evidence. */
|
||||
export function formatAgg(agg: AggregateRow): string {
|
||||
return `count=${agg.count} distinct_scopes=${agg.distinct_scopes} confidence=${agg.confidence.toFixed(2)} seen=[${agg.first_seen.slice(0, 10)}..${agg.last_seen.slice(0, 10)}]`;
|
||||
}
|
||||
@@ -1,269 +0,0 @@
|
||||
// kb_stats — on-demand dashboard numbers from the KB scratchpad
|
||||
// files. Reads data/_auditor/verdicts/*, data/_kb/audit_lessons.jsonl,
|
||||
// data/_kb/audit_facts.jsonl, data/_kb/audit_discrepancies.jsonl,
|
||||
// data/_kb/scrum_reviews.jsonl and prints:
|
||||
//
|
||||
// - verdict flip-flop rate (same SHA re-audited, verdict changed?)
|
||||
// - consensus discrepancy rate (N runs disagreed on a claim)
|
||||
// - confidence distribution from kb_index aggregator
|
||||
// - top N recurring entities from audit_facts
|
||||
// - fact growth over time
|
||||
// - scrum vs inference KB split
|
||||
//
|
||||
// Run: bun run auditor/kb_stats.ts
|
||||
// bun run auditor/kb_stats.ts --top 15 # show top 15 entities
|
||||
// bun run auditor/kb_stats.ts --json # machine-readable
|
||||
//
|
||||
// This is the "dashboard" without running Grafana. If someone really
|
||||
// wants a dashboard, wire this output into a static HTML page + cron.
|
||||
|
||||
import { readFile, readdir } from "node:fs/promises";
|
||||
import { join } from "node:path";
|
||||
import { aggregate } from "./kb_index.ts";
|
||||
|
||||
const REPO = "/home/profit/lakehouse";
|
||||
const VERDICTS_DIR = `${REPO}/data/_auditor/verdicts`;
|
||||
const AUDIT_LESSONS = `${REPO}/data/_kb/audit_lessons.jsonl`;
|
||||
const AUDIT_FACTS = `${REPO}/data/_kb/audit_facts.jsonl`;
|
||||
const AUDIT_DISCREPANCIES = `${REPO}/data/_kb/audit_discrepancies.jsonl`;
|
||||
const SCRUM_REVIEWS = `${REPO}/data/_kb/scrum_reviews.jsonl`;
|
||||
|
||||
interface Args {
|
||||
top: number;
|
||||
json: boolean;
|
||||
}
|
||||
|
||||
function parseArgs(argv: string[]): Args {
|
||||
const a: Args = { top: 10, json: false };
|
||||
for (let i = 2; i < argv.length; i++) {
|
||||
if (argv[i] === "--top") a.top = Number(argv[++i] ?? 10);
|
||||
else if (argv[i] === "--json") a.json = true;
|
||||
}
|
||||
return a;
|
||||
}
|
||||
|
||||
async function readJsonl<T = any>(path: string): Promise<T[]> {
|
||||
try {
|
||||
const raw = await readFile(path, "utf8");
|
||||
return raw.split("\n").filter(l => l.length > 0).map(l => {
|
||||
try { return JSON.parse(l) as T; } catch { return null as any; }
|
||||
}).filter(r => r !== null);
|
||||
} catch { return []; }
|
||||
}
|
||||
|
||||
async function loadVerdicts(): Promise<Array<{ pr: number; sha: string; overall: string; findings_total: number; findings_block: number; findings_warn: number }>> {
|
||||
let files: string[] = [];
|
||||
try { files = await readdir(VERDICTS_DIR); } catch { return []; }
|
||||
const out = [];
|
||||
for (const f of files) {
|
||||
if (!f.endsWith(".json")) continue;
|
||||
const m = f.match(/^(\d+)-([0-9a-f]+)\.json$/);
|
||||
if (!m) continue;
|
||||
try {
|
||||
const v = JSON.parse(await readFile(join(VERDICTS_DIR, f), "utf8"));
|
||||
out.push({
|
||||
pr: Number(m[1]),
|
||||
sha: m[2],
|
||||
overall: String(v.overall),
|
||||
findings_total: Number(v.metrics?.findings_total ?? 0),
|
||||
findings_block: Number(v.metrics?.findings_block ?? 0),
|
||||
findings_warn: Number(v.metrics?.findings_warn ?? 0),
|
||||
});
|
||||
} catch { /* skip corrupt */ }
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
interface Stats {
|
||||
audit_count: number;
|
||||
verdict_distribution: Record<string, number>;
|
||||
// Same PR with multiple SHAs — if verdicts differ, that's drift across
|
||||
// the PR's commit history. Not a flip-flop in the classical sense,
|
||||
// but worth surfacing (e.g. "PR #8 was block block req req block").
|
||||
per_pr_verdict_sequences: Record<number, string[]>;
|
||||
// For each PR with ≥ 2 audits, how many distinct verdicts did it
|
||||
// produce? 1 = stable; 2+ = some flipping.
|
||||
verdict_instability: { pr_count: number; pr_with_multiple_verdicts: number; pr_with_3plus: number };
|
||||
consensus: { discrepancy_count: number; tiebreaker_used: number; unresolved: number };
|
||||
kb: {
|
||||
audit_lessons_rows: number;
|
||||
audit_facts_rows: number;
|
||||
scrum_reviews_rows: number;
|
||||
distinct_finding_signatures: number;
|
||||
distinct_entities_across_prs: number;
|
||||
entities_in_2plus_prs: number;
|
||||
entities_in_5plus_prs: number;
|
||||
};
|
||||
fact_quality: {
|
||||
verifier_verdict_distribution: Record<string, number>;
|
||||
facts_dropped_by_verifier_total: number;
|
||||
extraction_success_rate: number;
|
||||
};
|
||||
top_entities: Array<{ name: string; distinct_prs: number; count: number; types: string[] }>;
|
||||
kb_by_source: Record<string, number>;
|
||||
}
|
||||
|
||||
async function collect(args: Args): Promise<Stats> {
|
||||
const verdicts = await loadVerdicts();
|
||||
const lessons = await readJsonl<any>(AUDIT_LESSONS);
|
||||
const facts = await readJsonl<any>(AUDIT_FACTS);
|
||||
const disc = await readJsonl<any>(AUDIT_DISCREPANCIES);
|
||||
const reviews = await readJsonl<any>(SCRUM_REVIEWS);
|
||||
|
||||
// Verdict stability
|
||||
const byPr: Record<number, string[]> = {};
|
||||
const verdictDist: Record<string, number> = {};
|
||||
for (const v of verdicts) {
|
||||
(byPr[v.pr] ??= []).push(v.overall);
|
||||
verdictDist[v.overall] = (verdictDist[v.overall] ?? 0) + 1;
|
||||
}
|
||||
let multi = 0, tri = 0;
|
||||
for (const [_, seq] of Object.entries(byPr)) {
|
||||
const distinct = new Set(seq);
|
||||
if (distinct.size >= 2) multi++;
|
||||
if (distinct.size >= 3) tri++;
|
||||
}
|
||||
|
||||
// Consensus drift
|
||||
const consensus = {
|
||||
discrepancy_count: disc.length,
|
||||
tiebreaker_used: disc.filter(d => String(d.resolution).startsWith("tiebreaker")).length,
|
||||
unresolved: disc.filter(d => d.resolution === "unresolved").length,
|
||||
};
|
||||
|
||||
// Lesson signatures
|
||||
const lessonAgg = await aggregate<any>(AUDIT_LESSONS, {
|
||||
keyFn: r => r?.signature,
|
||||
scopeFn: r => (r?.pr_number !== undefined ? `pr-${r.pr_number}` : undefined),
|
||||
});
|
||||
|
||||
// Entity aggregation across audit_facts rows
|
||||
interface EntAgg { distinct_prs: Set<number>; count: number; types: Set<string>; name: string; sources: Set<string> }
|
||||
const entAgg = new Map<string, EntAgg>();
|
||||
const sourceCount: Record<string, number> = {};
|
||||
let totalVerdictDist: Record<string, number> = { CORRECT: 0, INCORRECT: 0, UNVERIFIABLE: 0, UNCHECKED: 0 };
|
||||
let factsDroppedTotal = 0;
|
||||
let extractionsWithFacts = 0;
|
||||
|
||||
for (const row of facts) {
|
||||
const src = String(row.source ?? "unknown");
|
||||
sourceCount[src] = (sourceCount[src] ?? 0) + 1;
|
||||
const pr = Number(row.pr_number);
|
||||
if (Array.isArray(row.verifier_verdicts)) {
|
||||
for (const v of row.verifier_verdicts) {
|
||||
totalVerdictDist[v] = (totalVerdictDist[v] ?? 0) + 1;
|
||||
}
|
||||
}
|
||||
factsDroppedTotal += Number(row.facts_dropped_by_verifier ?? 0);
|
||||
if ((Array.isArray(row.facts) && row.facts.length > 0) || (Array.isArray(row.entities) && row.entities.length > 0)) {
|
||||
extractionsWithFacts++;
|
||||
}
|
||||
for (const e of Array.isArray(row.entities) ? row.entities : []) {
|
||||
const name = String(e?.name ?? "").trim();
|
||||
if (name.length < 3) continue;
|
||||
const key = name.toLowerCase();
|
||||
const agg = entAgg.get(key) ?? { distinct_prs: new Set(), count: 0, types: new Set(), name, sources: new Set() };
|
||||
agg.count++;
|
||||
if (Number.isFinite(pr) && pr > 0) agg.distinct_prs.add(pr);
|
||||
if (e?.type) agg.types.add(String(e.type));
|
||||
agg.sources.add(src);
|
||||
entAgg.set(key, agg);
|
||||
}
|
||||
}
|
||||
|
||||
const entitiesIn2Plus = Array.from(entAgg.values()).filter(a => a.distinct_prs.size >= 2).length;
|
||||
const entitiesIn5Plus = Array.from(entAgg.values()).filter(a => a.distinct_prs.size >= 5).length;
|
||||
const topEntities = Array.from(entAgg.values())
|
||||
.sort((a, b) => b.distinct_prs.size - a.distinct_prs.size || b.count - a.count)
|
||||
.slice(0, args.top)
|
||||
.map(a => ({
|
||||
name: a.name,
|
||||
distinct_prs: a.distinct_prs.size,
|
||||
count: a.count,
|
||||
types: Array.from(a.types),
|
||||
}));
|
||||
|
||||
const stats: Stats = {
|
||||
audit_count: verdicts.length,
|
||||
verdict_distribution: verdictDist,
|
||||
per_pr_verdict_sequences: byPr,
|
||||
verdict_instability: {
|
||||
pr_count: Object.keys(byPr).length,
|
||||
pr_with_multiple_verdicts: multi,
|
||||
pr_with_3plus: tri,
|
||||
},
|
||||
consensus,
|
||||
kb: {
|
||||
audit_lessons_rows: lessons.length,
|
||||
audit_facts_rows: facts.length,
|
||||
scrum_reviews_rows: reviews.length,
|
||||
distinct_finding_signatures: lessonAgg.size,
|
||||
distinct_entities_across_prs: entAgg.size,
|
||||
entities_in_2plus_prs: entitiesIn2Plus,
|
||||
entities_in_5plus_prs: entitiesIn5Plus,
|
||||
},
|
||||
fact_quality: {
|
||||
verifier_verdict_distribution: totalVerdictDist,
|
||||
facts_dropped_by_verifier_total: factsDroppedTotal,
|
||||
extraction_success_rate: facts.length > 0 ? extractionsWithFacts / facts.length : 0,
|
||||
},
|
||||
top_entities: topEntities,
|
||||
kb_by_source: sourceCount,
|
||||
};
|
||||
return stats;
|
||||
}
|
||||
|
||||
function renderHuman(s: Stats): string {
|
||||
const lines: string[] = [];
|
||||
lines.push("═══ KB STATS ═══");
|
||||
lines.push("");
|
||||
lines.push(`Audits: ${s.audit_count} total across ${s.verdict_instability.pr_count} distinct PRs`);
|
||||
lines.push(`Verdicts: ${Object.entries(s.verdict_distribution).map(([k, v]) => `${k}=${v}`).join(" ")}`);
|
||||
const multiplePct = s.verdict_instability.pr_count > 0
|
||||
? Math.round(100 * s.verdict_instability.pr_with_multiple_verdicts / s.verdict_instability.pr_count)
|
||||
: 0;
|
||||
lines.push(`Verdict instability: ${s.verdict_instability.pr_with_multiple_verdicts}/${s.verdict_instability.pr_count} PRs had 2+ distinct verdicts (${multiplePct}%) — 3+ distinct: ${s.verdict_instability.pr_with_3plus}`);
|
||||
lines.push("");
|
||||
lines.push("─── Consensus ───");
|
||||
lines.push(` discrepancies logged: ${s.consensus.discrepancy_count}`);
|
||||
lines.push(` tiebreaker used: ${s.consensus.tiebreaker_used}`);
|
||||
lines.push(` unresolved: ${s.consensus.unresolved}`);
|
||||
const dRate = s.audit_count > 0 ? (100 * s.consensus.discrepancy_count / s.audit_count).toFixed(1) : "0";
|
||||
lines.push(` discrepancy rate: ${dRate}% of audits`);
|
||||
lines.push("");
|
||||
lines.push("─── KB size ───");
|
||||
lines.push(` audit_lessons.jsonl: ${s.kb.audit_lessons_rows} rows, ${s.kb.distinct_finding_signatures} distinct signatures`);
|
||||
lines.push(` audit_facts.jsonl: ${s.kb.audit_facts_rows} rows, ${s.kb.distinct_entities_across_prs} distinct entities`);
|
||||
lines.push(` scrum_reviews.jsonl: ${s.kb.scrum_reviews_rows} rows`);
|
||||
lines.push(` entities in 2+ PRs: ${s.kb.entities_in_2plus_prs}`);
|
||||
lines.push(` entities in 5+ PRs: ${s.kb.entities_in_5plus_prs} ← strong cross-cutting signal`);
|
||||
lines.push("");
|
||||
lines.push("─── Fact quality ───");
|
||||
const v = s.fact_quality.verifier_verdict_distribution;
|
||||
lines.push(` verifier verdicts: CORRECT=${v.CORRECT ?? 0} UNVERIFIABLE=${v.UNVERIFIABLE ?? 0} UNCHECKED=${v.UNCHECKED ?? 0} INCORRECT=${v.INCORRECT ?? 0}`);
|
||||
lines.push(` facts dropped by verifier: ${s.fact_quality.facts_dropped_by_verifier_total}`);
|
||||
lines.push(` extraction success rate: ${(s.fact_quality.extraction_success_rate * 100).toFixed(1)}%`);
|
||||
lines.push("");
|
||||
lines.push("─── KB sources ───");
|
||||
for (const [src, n] of Object.entries(s.kb_by_source)) {
|
||||
lines.push(` ${src}: ${n}`);
|
||||
}
|
||||
lines.push("");
|
||||
lines.push(`─── Top ${s.top_entities.length} recurring entities ───`);
|
||||
for (const e of s.top_entities) {
|
||||
lines.push(` [${e.distinct_prs} PRs × ${e.count} obs] ${e.name} (${e.types.join(",")})`);
|
||||
}
|
||||
return lines.join("\n");
|
||||
}
|
||||
|
||||
async function main() {
|
||||
const args = parseArgs(process.argv);
|
||||
const stats = await collect(args);
|
||||
if (args.json) {
|
||||
console.log(JSON.stringify(stats, (_, v) => v instanceof Set ? Array.from(v) : v, 2));
|
||||
} else {
|
||||
console.log(renderHuman(stats));
|
||||
}
|
||||
}
|
||||
|
||||
main().catch(e => { console.error("[kb_stats] fatal:", e); process.exit(1); });
|
||||
@@ -1,62 +0,0 @@
|
||||
// ═══════════════════════════════════════════════════════════════════
|
||||
// YOU WRITE THIS FILE. Policy decides what blocks vs what's a comment.
|
||||
// Defaults are opinionated on the "stop clicking past placeholder"
|
||||
// side — easier to loosen than to tighten when you're watching the
|
||||
// auditor behave in live PRs.
|
||||
// ═══════════════════════════════════════════════════════════════════
|
||||
|
||||
import type { Finding, Verdict } from "./types.ts";
|
||||
|
||||
/// Translate the four-check output into a single verdict. This is the
|
||||
/// single pane of glass the auditor operates on — tune thresholds here.
|
||||
export function assembleVerdict(
|
||||
findings: Finding[],
|
||||
metrics: Record<string, number>,
|
||||
pr_number: number,
|
||||
head_sha: string,
|
||||
): Verdict {
|
||||
const blocking = findings.filter(f => f.severity === "block");
|
||||
const warning = findings.filter(f => f.severity === "warn");
|
||||
|
||||
let overall: Verdict["overall"];
|
||||
let one_liner: string;
|
||||
|
||||
if (blocking.length > 0) {
|
||||
overall = "block";
|
||||
one_liner = `${blocking.length} blocking issue${blocking.length > 1 ? "s" : ""}: ${blocking[0].summary}`;
|
||||
} else if (warning.length >= 3) {
|
||||
// Three or more warnings is a block — death by a thousand cuts.
|
||||
overall = "request_changes";
|
||||
one_liner = `${warning.length} warnings — see review`;
|
||||
} else if (warning.length > 0) {
|
||||
overall = "request_changes";
|
||||
one_liner = warning[0].summary;
|
||||
} else {
|
||||
overall = "approve";
|
||||
one_liner = `all checks passed (${findings.length} findings, all info)`;
|
||||
}
|
||||
|
||||
return {
|
||||
pr_number,
|
||||
head_sha,
|
||||
audited_at: new Date().toISOString(),
|
||||
overall,
|
||||
findings,
|
||||
metrics,
|
||||
one_liner,
|
||||
};
|
||||
}
|
||||
|
||||
/// Which strength-of-claim warrants which severity when evidence is
|
||||
/// weak? A "Phase X shipped" claim with zero integration tests is a
|
||||
/// blocker. A "should work" claim with no test is a warn.
|
||||
export function severityFromClaimEvidence(
|
||||
claim_strength: "weak" | "moderate" | "strong",
|
||||
evidence_grade: "none" | "partial" | "full",
|
||||
): "info" | "warn" | "block" {
|
||||
if (evidence_grade === "full") return "info";
|
||||
if (claim_strength === "strong" && evidence_grade === "none") return "block";
|
||||
if (claim_strength === "strong" && evidence_grade === "partial") return "warn";
|
||||
if (claim_strength === "moderate" && evidence_grade === "none") return "warn";
|
||||
return "info";
|
||||
}
|
||||
@@ -1,72 +0,0 @@
|
||||
// Shared types for the claim-auditor. Every field exists for a reason;
|
||||
// if something can't be verified from a check, it goes into `evidence`
|
||||
// so the verdict is inspectable, not a black box.
|
||||
|
||||
export type CheckKind = "static" | "dynamic" | "inference" | "kb_query";
|
||||
|
||||
export type Severity = "info" | "warn" | "block";
|
||||
|
||||
export interface Claim {
|
||||
// Verbatim phrase that raised the claim — e.g. "Phase 38 shipped",
|
||||
// "verified end-to-end", "works after restart". Used as the "what
|
||||
// does the author assert" input to downstream checks.
|
||||
text: string;
|
||||
// Where it came from. `commit_sha` is the short hash; `location`
|
||||
// is a file:line for in-diff claims, or "pr_body" / "commit_message".
|
||||
commit_sha: string;
|
||||
location: string;
|
||||
// Heuristic rating of how strong the claim is. "green+tested"
|
||||
// is strong; "should work" is weak. Drives sensitivity — stronger
|
||||
// claims get harder-blocked on weak evidence.
|
||||
//
|
||||
// "empirical" is a separate class: runtime measurements like
|
||||
// "N cloud calls" / "306s end-to-end" / "accepted on attempt N".
|
||||
// These cannot be verified from a static diff — only from the test
|
||||
// output that produced them. Inference skips diff-verification for
|
||||
// empirical claims; they become info-level context unless a future
|
||||
// runtime_evidence check contradicts them.
|
||||
strength: "weak" | "moderate" | "strong" | "empirical";
|
||||
}
|
||||
|
||||
export interface Finding {
|
||||
check: CheckKind;
|
||||
severity: Severity;
|
||||
claim_text?: string;
|
||||
// Free-form short description: "field added but never read", "no
|
||||
// test covers this code path", "cloud model says placeholder".
|
||||
summary: string;
|
||||
// Concrete evidence: file paths, line numbers, log excerpts, test
|
||||
// output, cloud-model verdict. No handwaving.
|
||||
evidence: string[];
|
||||
}
|
||||
|
||||
export interface Verdict {
|
||||
pr_number: number;
|
||||
head_sha: string;
|
||||
audited_at: string;
|
||||
overall: "approve" | "request_changes" | "block";
|
||||
findings: Finding[];
|
||||
// Real numbers that downstream policy can gate on. e.g. if the
|
||||
// hybrid test produced latency numbers or token counts, they
|
||||
// surface here so /auditor/history is queryable.
|
||||
metrics: Record<string, number>;
|
||||
// Short one-line justification for the `overall` verdict. What
|
||||
// gets posted as the commit-status description in Gitea (max 140
|
||||
// chars) must fit here.
|
||||
one_liner: string;
|
||||
}
|
||||
|
||||
export interface PrSnapshot {
|
||||
number: number;
|
||||
head_sha: string;
|
||||
base_sha: string;
|
||||
title: string;
|
||||
body: string;
|
||||
state: "open" | "closed" | "merged";
|
||||
author: string;
|
||||
// Array of commit messages in the PR (not diffs — those are
|
||||
// fetched on-demand per-check).
|
||||
commits: Array<{ sha: string; message: string; author: string }>;
|
||||
// File paths touched by the PR, with lines-added / lines-removed.
|
||||
files: Array<{ path: string; additions: number; deletions: number }>;
|
||||
}
|
||||
-127
@@ -1,127 +0,0 @@
|
||||
# Lakehouse PR-Bot
|
||||
|
||||
A local sub-agent that reads the PRD, asks a cloud model for a small change
|
||||
proposal, applies it, runs tests, and opens a draft PR on Gitea. Manual-only
|
||||
right now — no systemd, no cron. Run one cycle at a time and watch.
|
||||
|
||||
## Run one cycle
|
||||
|
||||
```bash
|
||||
cd /home/profit/lakehouse
|
||||
bun run bot/cycle.ts
|
||||
```
|
||||
|
||||
Prerequisites:
|
||||
- Working tree must be clean (`git status` shows nothing). The bot refuses on a dirty tree so its changes don't mix with in-flight work.
|
||||
- Sidecar running on :3200, gateway on :3100, observer on :3800.
|
||||
- At least one PRD line tagged `[bot-eligible]` (see below).
|
||||
- Gitea PAT configured in `~/.git-credentials` (already set up).
|
||||
|
||||
## Tag a gap as bot-eligible
|
||||
|
||||
Add `[bot-eligible]` to any PRD line the bot is allowed to work on:
|
||||
|
||||
```md
|
||||
- [ ] Add a unit test for parse_city_state covering "South Bend, IN" edge case [bot-eligible]
|
||||
```
|
||||
|
||||
The bot scans `docs/PRD.md` for these tags. Each tagged line becomes a candidate.
|
||||
Start small — one tag at a time — until you trust the loop.
|
||||
|
||||
## Stop it mid-cycle
|
||||
|
||||
Create the pause file:
|
||||
|
||||
```bash
|
||||
touch /home/profit/lakehouse/bot.paused
|
||||
```
|
||||
|
||||
The next `bot/cycle.ts` invocation exits immediately with `skipped_pause`.
|
||||
(It does NOT kill a cycle already in-flight — use `Ctrl-C` or `pkill -f bot/cycle.ts` for that.)
|
||||
|
||||
## Budget
|
||||
|
||||
- **20 cloud calls/day, 160k tokens/day** (hard ceiling, see `bot/cost.ts`).
|
||||
- Tracked in `data/_bot/cost-YYYY-MM-DD.json`.
|
||||
- Resets at UTC midnight.
|
||||
|
||||
## Cycle outcomes
|
||||
|
||||
Every run writes a result to `data/_bot/cycles/{cycle_id}.json` and POSTs an event to the observer on :3800.
|
||||
Possible outcomes:
|
||||
|
||||
| Outcome | Meaning |
|
||||
|---|---|
|
||||
| `ok` | PR opened |
|
||||
| `cycle_noop` | proposal applied, every file was **identical to current content** (Mem0 NOOP); no test run, no PR |
|
||||
| `skipped_pause` | pause file present |
|
||||
| `skipped_cost` | daily budget exhausted |
|
||||
| `skipped_policy` | `policy.shouldRunCycle` said no |
|
||||
| `skipped_dirty_tree` | uncommitted changes present |
|
||||
| `skipped_no_gap` | no `[bot-eligible]` tags in PRD |
|
||||
| `model_failed` | sidecar/cloud model errored or returned unparseable JSON |
|
||||
| `proposal_rejected` | `policy.scoreProposal` rejected it (size, path, etc.) |
|
||||
| `apply_failed` | file write errored |
|
||||
| `tests_failed` | cargo or bun test red |
|
||||
| `pr_skipped_by_policy` | green tests, but `policy.shouldOpenPR` said no |
|
||||
| `pr_failed` | Gitea API or git push errored |
|
||||
|
||||
## Mem0-aligned apply semantics
|
||||
|
||||
`apply.ts` categorizes every proposed file into one of three modes:
|
||||
|
||||
| Mode | Trigger | Action |
|
||||
|---|---|---|
|
||||
| **ADD** | `is_new: true` and file doesn't exist | Create the file |
|
||||
| **UPDATE** | `is_new: false`, file exists, content **differs** from current | Overwrite |
|
||||
| **NOOP** | `is_new: false`, file exists, content **matches** current exactly | Skip (no write, no diff) |
|
||||
|
||||
If every file in the proposal is NOOP, the cycle short-circuits to `cycle_noop` **before** running tests or opening a PR. Mismatched shapes (`is_new:true` but file exists, or `is_new:false` but file missing) become `apply_failed` — model state confusion is surfaced, not papered over.
|
||||
|
||||
PR bodies report the three counts separately so reviewers see what actually changed vs. what was confirmed identical.
|
||||
|
||||
## How the bot compounds over cycles
|
||||
|
||||
Every finished cycle persists a `CycleResult` to `data/_bot/cycles/{cycle_id}.json`. At the **start** of the next cycle, `bot/kb.ts::loadHistory(gap_id)` scans that directory, filters to prior cycles on the same gap, and returns the five most recent outcomes (`ok`, `tests_failed`, `proposal_rejected`, `cycle_noop`, etc.) with their reasons and touched files.
|
||||
|
||||
Those outcomes are summarized into a compact block and **injected into the cloud prompt** before asking for a new proposal. The model sees things like:
|
||||
|
||||
```
|
||||
Prior attempts on this gap (3 most recent):
|
||||
- 2026-04-22 03:15 UTC — tests_failed
|
||||
reason: cargo test -p vectord::lance failed on field_type_coerce
|
||||
- 2026-04-22 02:48 UTC — proposal_rejected
|
||||
reason: touched forbidden path (docs/ADR-): docs/ADR-019-update.md
|
||||
- 2026-04-22 01:23 UTC — ok PR: https://git.agentview.dev/profit/lakehouse/pulls/142
|
||||
reason: PR #142 opened
|
||||
|
||||
Learn from these: build on what worked, avoid paths that failed.
|
||||
```
|
||||
|
||||
This is the same compounding pattern `scenario.ts` uses via `kb.loadRecommendation` — the bot's cycles are the bot's memory. No embedding, no separate jsonl, no cross-run orchestration required. First cycle on a new gap skips the block cleanly.
|
||||
|
||||
The observer events POSTed on every cycle carry the same data, so `GET :3800/stats?source=bot` aggregates "% of bot cycles on similar gaps that landed a PR" without extra plumbing.
|
||||
|
||||
## Where YOU edit
|
||||
|
||||
`bot/policy.ts` — four small functions that define what the bot does. The rest
|
||||
(`propose.ts`, `apply.ts`, `test.ts`, `pr.ts`, `cycle.ts`, `kb.ts`) is mechanical
|
||||
orchestration — you shouldn't need to touch it unless a pipeline changes.
|
||||
|
||||
One policy upgrade the history opens up: in `shouldRunCycle`, you can now bail
|
||||
out if the same gap has failed N times in a row. Example addition:
|
||||
|
||||
```ts
|
||||
import { loadHistory, statsFor } from "./kb.ts";
|
||||
// inside shouldRunCycle — but you'd need the gap, which is picked later.
|
||||
// A more natural place is scoreProposal: if prior N failures on this
|
||||
// gap's path+summary pattern, reject before testing.
|
||||
```
|
||||
|
||||
## Hard-coded guardrails (not in policy — can't be disabled)
|
||||
|
||||
- Bot **never deletes files** (`apply.ts` has no delete path).
|
||||
- Bot **never touches** `.git/`, `secrets`, `lakehouse.toml`, `docs/ADR-*`, `docs/DECISIONS.md`, `docs/PRD.md`, `/etc/`, `/root/`, `Cargo.lock`.
|
||||
- All paths validated for traversal (`path/../`) and repo-escape.
|
||||
- PRs always open as **draft** — never auto-merge.
|
||||
- Budget check is before the policy function runs — no way to override the daily cap from `policy.ts`.
|
||||
@@ -1,86 +0,0 @@
|
||||
// Apply a proposal to disk with Mem0-aligned ADD / UPDATE / NOOP
|
||||
// semantics. Whole-file writes only — never deletes. All paths
|
||||
// validated repo-relative + path-traversal-free.
|
||||
//
|
||||
// Three outcomes per file:
|
||||
// ADD — file didn't exist, new content written (is_new must be true)
|
||||
// UPDATE — file existed, content differs from current, overwritten
|
||||
// NOOP — file existed, content identical to current, nothing written
|
||||
//
|
||||
// The NOOP case is important: a cloud proposal that "re-derives" an
|
||||
// existing file shouldn't churn the repo or burn test cycles.
|
||||
|
||||
import { writeFile, readFile, mkdir, access } from "node:fs/promises";
|
||||
import { dirname, resolve, relative } from "node:path";
|
||||
import type { ApplyOutcome, Proposal } from "./types.ts";
|
||||
|
||||
const REPO_ROOT = "/home/profit/lakehouse";
|
||||
|
||||
export async function applyProposal(p: Proposal): Promise<ApplyOutcome> {
|
||||
const added: string[] = [];
|
||||
const updated: string[] = [];
|
||||
const noop: string[] = [];
|
||||
const errors: string[] = [];
|
||||
|
||||
for (const f of p.files) {
|
||||
const err = validatePath(f.path);
|
||||
if (err) {
|
||||
errors.push(`${f.path}: ${err}`);
|
||||
continue;
|
||||
}
|
||||
const abs = resolve(REPO_ROOT, f.path);
|
||||
try {
|
||||
const exists = await fileExists(abs);
|
||||
|
||||
if (f.is_new && exists) {
|
||||
errors.push(`${f.path}: is_new=true but file exists`);
|
||||
continue;
|
||||
}
|
||||
if (!f.is_new && !exists) {
|
||||
// Proposal claims to update a file that doesn't exist.
|
||||
// Refuse rather than silently ADD — model is confused about
|
||||
// repo state, and we want that surfaced, not papered over.
|
||||
errors.push(`${f.path}: is_new=false but file does not exist`);
|
||||
continue;
|
||||
}
|
||||
|
||||
if (!f.is_new) {
|
||||
// UPDATE candidate — compare to detect NOOP.
|
||||
const current = await readFile(abs, "utf8");
|
||||
if (current === f.content) {
|
||||
noop.push(f.path);
|
||||
continue;
|
||||
}
|
||||
await writeFile(abs, f.content);
|
||||
updated.push(f.path);
|
||||
} else {
|
||||
// ADD — file doesn't exist yet.
|
||||
await mkdir(dirname(abs), { recursive: true });
|
||||
await writeFile(abs, f.content);
|
||||
added.push(f.path);
|
||||
}
|
||||
} catch (e) {
|
||||
errors.push(`${f.path}: ${(e as Error).message}`);
|
||||
}
|
||||
}
|
||||
return { added, updated, noop, errors };
|
||||
}
|
||||
|
||||
function validatePath(p: string): string | null {
|
||||
if (!p) return "empty path";
|
||||
if (p.startsWith("/")) return "must be repo-relative";
|
||||
if (p.startsWith("..") || p.includes("/../")) return "path traversal";
|
||||
const abs = resolve(REPO_ROOT, p);
|
||||
const rel = relative(REPO_ROOT, abs);
|
||||
if (rel.startsWith("..")) return "resolves outside repo";
|
||||
return null;
|
||||
}
|
||||
|
||||
async function fileExists(abs: string): Promise<boolean> {
|
||||
try {
|
||||
await access(abs);
|
||||
return true;
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
-52
@@ -1,52 +0,0 @@
|
||||
// Daily cost tracker. Resets at UTC midnight by keying the file by date.
|
||||
// Hard ceilings are the policy surface — the bot refuses to call the
|
||||
// cloud once the budget is exhausted for the day.
|
||||
|
||||
import { readFile, writeFile, mkdir } from "node:fs/promises";
|
||||
import { join } from "node:path";
|
||||
import type { CostState } from "./types.ts";
|
||||
|
||||
const COST_DIR = "/home/profit/lakehouse/data/_bot";
|
||||
|
||||
export const DAILY_CALLS_BUDGET = 20;
|
||||
export const DAILY_TOKENS_BUDGET = 8000 * DAILY_CALLS_BUDGET; // 160k tokens/day ceiling
|
||||
|
||||
function todayUtc(): string {
|
||||
return new Date().toISOString().slice(0, 10);
|
||||
}
|
||||
|
||||
function costFile(date: string): string {
|
||||
return join(COST_DIR, `cost-${date}.json`);
|
||||
}
|
||||
|
||||
export async function readCost(): Promise<CostState> {
|
||||
const date = todayUtc();
|
||||
try {
|
||||
const raw = await readFile(costFile(date), "utf8");
|
||||
return JSON.parse(raw) as CostState;
|
||||
} catch {
|
||||
return { date, calls: 0, tokens: 0 };
|
||||
}
|
||||
}
|
||||
|
||||
export async function recordCost(calls: number, tokens: number): Promise<CostState> {
|
||||
await mkdir(COST_DIR, { recursive: true });
|
||||
const current = await readCost();
|
||||
const updated: CostState = {
|
||||
date: current.date,
|
||||
calls: current.calls + calls,
|
||||
tokens: current.tokens + tokens,
|
||||
};
|
||||
await writeFile(costFile(updated.date), JSON.stringify(updated, null, 2));
|
||||
return updated;
|
||||
}
|
||||
|
||||
export function budgetCheck(c: CostState): { ok: boolean; reason: string } {
|
||||
if (c.calls >= DAILY_CALLS_BUDGET) {
|
||||
return { ok: false, reason: `daily call budget exhausted (${c.calls}/${DAILY_CALLS_BUDGET})` };
|
||||
}
|
||||
if (c.tokens >= DAILY_TOKENS_BUDGET) {
|
||||
return { ok: false, reason: `daily token budget exhausted (${c.tokens}/${DAILY_TOKENS_BUDGET})` };
|
||||
}
|
||||
return { ok: true, reason: `${c.calls}/${DAILY_CALLS_BUDGET} calls, ${c.tokens}/${DAILY_TOKENS_BUDGET} tokens used today` };
|
||||
}
|
||||
-345
@@ -1,345 +0,0 @@
|
||||
// Single-cycle orchestrator. This is the manual entry point:
|
||||
// bun run bot/cycle.ts
|
||||
//
|
||||
// Steps (each step can short-circuit with a clear outcome):
|
||||
// 1. Pause-file check → skipped_pause
|
||||
// 2. Cost budget check → skipped_cost
|
||||
// 3. Gather CycleContext (dirty tree, last cycle, autotune status)
|
||||
// 4. policy.shouldRunCycle → skipped_policy
|
||||
// 5. findGaps + policy.pickGap → skipped_no_gap
|
||||
// 6. generateProposal → model_failed
|
||||
// 7. policy.scoreProposal → proposal_rejected
|
||||
// 8. git checkout new branch
|
||||
// 9. applyProposal → apply_failed
|
||||
// 10. runTests → tests_failed
|
||||
// 11. policy.shouldOpenPR → pr_skipped_by_policy
|
||||
// 12. git commit + push
|
||||
// 13. openPr → pr_failed / ok
|
||||
// 14. recordCost + postEvent (every outcome)
|
||||
// 15. write cycle result to data/_bot/cycles/{cycle_id}.json
|
||||
|
||||
import { readFile, writeFile, mkdir, access } from "node:fs/promises";
|
||||
import { createHash, randomUUID } from "node:crypto";
|
||||
import { spawnSync } from "node:child_process";
|
||||
import { join } from "node:path";
|
||||
|
||||
import type { CycleContext, CycleResult, CycleOutcome, Proposal } from "./types.ts";
|
||||
import * as policy from "./policy.ts";
|
||||
import { readCost, recordCost, budgetCheck, DAILY_CALLS_BUDGET, DAILY_TOKENS_BUDGET } from "./cost.ts";
|
||||
import { findGaps, generateProposal } from "./propose.ts";
|
||||
import { applyProposal } from "./apply.ts";
|
||||
import { runTests } from "./test.ts";
|
||||
import { openPr } from "./pr.ts";
|
||||
import { postEvent } from "./observer.ts";
|
||||
import { loadHistory, summarizeHistory, statsFor } from "./kb.ts";
|
||||
|
||||
const REPO_ROOT = "/home/profit/lakehouse";
|
||||
const PAUSE_FILE = `${REPO_ROOT}/bot.paused`;
|
||||
const CYCLES_DIR = `${REPO_ROOT}/data/_bot/cycles`;
|
||||
const GATEWAY_URL = process.env.LH_GATEWAY_URL ?? "http://localhost:3100";
|
||||
|
||||
function log(msg: string) { console.log(`[bot] ${msg}`); }
|
||||
|
||||
async function fileExists(p: string): Promise<boolean> {
|
||||
try { await access(p); return true; } catch { return false; }
|
||||
}
|
||||
|
||||
function git(args: string[]): { code: number; stdout: string; stderr: string } {
|
||||
const r = spawnSync("git", args, { cwd: REPO_ROOT, encoding: "utf8" });
|
||||
return { code: r.status ?? -1, stdout: r.stdout ?? "", stderr: r.stderr ?? "" };
|
||||
}
|
||||
|
||||
async function gatherContext(cost: { calls: number; tokens: number }): Promise<CycleContext> {
|
||||
// Last cycle: newest file in CYCLES_DIR.
|
||||
let lastCycleAt: string | null = null;
|
||||
let lastCycleGapId: string | null = null;
|
||||
try {
|
||||
const { readdir, stat } = await import("node:fs/promises");
|
||||
const entries = await readdir(CYCLES_DIR).catch(() => [] as string[]);
|
||||
let newestTs = 0;
|
||||
let newestPath: string | null = null;
|
||||
for (const e of entries) {
|
||||
if (!e.endsWith(".json")) continue;
|
||||
const p = join(CYCLES_DIR, e);
|
||||
const s = await stat(p);
|
||||
if (s.mtimeMs > newestTs) { newestTs = s.mtimeMs; newestPath = p; }
|
||||
}
|
||||
if (newestPath) {
|
||||
const r = JSON.parse(await readFile(newestPath, "utf8")) as CycleResult;
|
||||
lastCycleAt = r.started_at;
|
||||
lastCycleGapId = r.gap?.id ?? null;
|
||||
}
|
||||
} catch {}
|
||||
|
||||
// Autotune status — if the agent is busy, skip this cycle to avoid
|
||||
// contention on shared Ollama GPU memory.
|
||||
let autotuneBusy = false;
|
||||
try {
|
||||
const r = await fetch(`${GATEWAY_URL}/vectors/agent/status`, { signal: AbortSignal.timeout(2000) });
|
||||
if (r.ok) {
|
||||
const j = await r.json() as any;
|
||||
autotuneBusy = Boolean(j.running) && (j.queue_depth ?? 0) > 0;
|
||||
}
|
||||
} catch {}
|
||||
|
||||
// Dirty tree check.
|
||||
const status = git(["status", "--porcelain"]);
|
||||
const workingTreeDirty = status.code === 0 && status.stdout.trim().length > 0;
|
||||
|
||||
return {
|
||||
startedAt: new Date().toISOString(),
|
||||
dailyCallsUsed: cost.calls,
|
||||
dailyCallsBudget: DAILY_CALLS_BUDGET,
|
||||
dailyTokensUsed: cost.tokens,
|
||||
dailyTokensBudget: DAILY_TOKENS_BUDGET,
|
||||
lastCycleAt,
|
||||
lastCycleGapId,
|
||||
autotuneBusy,
|
||||
workingTreeDirty,
|
||||
};
|
||||
}
|
||||
|
||||
async function persistResult(res: CycleResult): Promise<void> {
|
||||
await mkdir(CYCLES_DIR, { recursive: true });
|
||||
await writeFile(join(CYCLES_DIR, `${res.cycle_id}.json`), JSON.stringify(res, null, 2));
|
||||
}
|
||||
|
||||
function sigHash(gap_id: string | null, proposalSummary: string): string {
|
||||
const h = createHash("sha256");
|
||||
h.update(gap_id ?? "no-gap");
|
||||
h.update("|");
|
||||
h.update(proposalSummary);
|
||||
return h.digest("hex").slice(0, 16);
|
||||
}
|
||||
|
||||
function emit(res: CycleResult): Promise<void> {
|
||||
return postEvent({
|
||||
source: "bot",
|
||||
cycle_id: res.cycle_id,
|
||||
sig_hash: sigHash(res.gap?.id ?? null, res.proposal?.summary ?? ""),
|
||||
event_kind: res.outcome,
|
||||
ok: res.outcome === "ok",
|
||||
staffer_id: res.proposal?.model_used,
|
||||
turns: res.cloud_calls,
|
||||
duration_ms: new Date(res.ended_at).getTime() - new Date(res.started_at).getTime(),
|
||||
extra: {
|
||||
reason: res.reason,
|
||||
pr_url: res.prUrl,
|
||||
files_added: res.filesAdded,
|
||||
files_updated: res.filesUpdated,
|
||||
files_noop: res.filesNoop,
|
||||
tests_green: res.testsGreen,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
async function main() {
|
||||
const cycle_id = `bot-${new Date().toISOString().replace(/[:.]/g, "-")}-${randomUUID().slice(0, 8)}`;
|
||||
const started_at = new Date().toISOString();
|
||||
log(`cycle ${cycle_id} starting`);
|
||||
|
||||
const finish = async (outcome: CycleOutcome, reason: string, extra: Partial<CycleResult> = {}): Promise<CycleResult> => {
|
||||
const res: CycleResult = {
|
||||
cycle_id,
|
||||
started_at,
|
||||
ended_at: new Date().toISOString(),
|
||||
outcome,
|
||||
reason,
|
||||
gap: null,
|
||||
proposal: null,
|
||||
filesAdded: [],
|
||||
filesUpdated: [],
|
||||
filesNoop: [],
|
||||
testsGreen: null,
|
||||
testsOutput: "",
|
||||
prUrl: null,
|
||||
tokens_used: 0,
|
||||
cloud_calls: 0,
|
||||
...extra,
|
||||
};
|
||||
await persistResult(res);
|
||||
await emit(res);
|
||||
log(`→ ${outcome}: ${reason}`);
|
||||
return res;
|
||||
};
|
||||
|
||||
// 1. Pause.
|
||||
if (await fileExists(PAUSE_FILE)) {
|
||||
return finish("skipped_pause", `pause file at ${PAUSE_FILE}`);
|
||||
}
|
||||
|
||||
// 2. Cost.
|
||||
const cost = await readCost();
|
||||
const bc = budgetCheck(cost);
|
||||
if (!bc.ok) return finish("skipped_cost", bc.reason);
|
||||
log(`cost: ${bc.reason}`);
|
||||
|
||||
// 3+4. Context + policy.shouldRunCycle.
|
||||
const ctx = await gatherContext(cost);
|
||||
const srq = policy.shouldRunCycle(ctx);
|
||||
if (!srq.run) {
|
||||
return finish(ctx.workingTreeDirty ? "skipped_dirty_tree" : "skipped_policy", srq.reason);
|
||||
}
|
||||
|
||||
// 5. Gaps + policy.pickGap.
|
||||
const gaps = await findGaps();
|
||||
log(`found ${gaps.length} [bot-eligible] gap(s) in PRD`);
|
||||
const gapPick = policy.pickGap(gaps);
|
||||
if (!gapPick.gap) return finish("skipped_no_gap", gapPick.reason);
|
||||
const gap = gapPick.gap;
|
||||
log(`picked gap ${gap.id}: "${gap.prd_line.slice(0, 80)}"`);
|
||||
|
||||
// 6. KB lookup + Proposal.
|
||||
// Bot cycles compound: prior cycles on this gap inform the current
|
||||
// proposal. Summary is a compact block injected into the cloud prompt
|
||||
// so the model can build on past successes and steer around past
|
||||
// failures. Empty string when this is the first cycle on this gap.
|
||||
const history = await loadHistory(gap.id, 5);
|
||||
const stats = statsFor(history);
|
||||
if (history.length > 0) {
|
||||
log(`kb: ${stats.attempts} prior cycle(s) — ${stats.pr_opened} pr / ${stats.tests_failed} tests_failed / ${stats.proposal_rejected} rejected / ${stats.noop} noop`);
|
||||
} else {
|
||||
log(`kb: first cycle on this gap`);
|
||||
}
|
||||
const historySummary = summarizeHistory(history);
|
||||
|
||||
let proposal: Proposal;
|
||||
try {
|
||||
proposal = await generateProposal(gap, historySummary);
|
||||
log(`proposal: ${proposal.summary} (${proposal.files.length} files, ~${proposal.estimated_loc} LOC, ${proposal.tokens_used} tokens)`);
|
||||
} catch (e) {
|
||||
await recordCost(1, 0);
|
||||
return finish("model_failed", (e as Error).message, { gap, cloud_calls: 1 });
|
||||
}
|
||||
await recordCost(1, proposal.tokens_used);
|
||||
|
||||
// 7. Score.
|
||||
const score = policy.scoreProposal(proposal);
|
||||
if (!score.accept) {
|
||||
return finish("proposal_rejected", score.reason, { gap, proposal, tokens_used: proposal.tokens_used, cloud_calls: 1 });
|
||||
}
|
||||
log(`proposal accepted: ${score.reason}`);
|
||||
|
||||
// 8. Branch.
|
||||
const dayStamp = new Date().toISOString().slice(0, 10).replace(/-/g, "");
|
||||
const branch = `bot/cycle-${dayStamp}-${gap.id}`;
|
||||
// Branch must be fresh. If one already exists from a prior failed cycle
|
||||
// for the same gap, bail — don't silently overwrite or reuse stale state.
|
||||
const check = git(["rev-parse", "--verify", `refs/heads/${branch}`]);
|
||||
if (check.code === 0) {
|
||||
return finish("apply_failed",
|
||||
`branch ${branch} already exists — delete it manually if you want a fresh cycle for this gap`,
|
||||
{ gap, proposal, tokens_used: proposal.tokens_used, cloud_calls: 1 });
|
||||
}
|
||||
const co = git(["checkout", "-b", branch]);
|
||||
if (co.code !== 0) {
|
||||
return finish("apply_failed", `git checkout -b failed: ${co.stderr}`,
|
||||
{ gap, proposal, tokens_used: proposal.tokens_used, cloud_calls: 1 });
|
||||
}
|
||||
log(`on branch ${branch}`);
|
||||
|
||||
// 9. Apply — Mem0 ADD/UPDATE/NOOP semantics.
|
||||
const ap = await applyProposal(proposal);
|
||||
if (ap.errors.length > 0) {
|
||||
git(["checkout", "main"]);
|
||||
git(["branch", "-D", branch]);
|
||||
return finish("apply_failed", ap.errors.join("; "),
|
||||
{ gap, proposal, tokens_used: proposal.tokens_used, cloud_calls: 1 });
|
||||
}
|
||||
const changed = [...ap.added, ...ap.updated];
|
||||
log(`apply: ${ap.added.length} add, ${ap.updated.length} update, ${ap.noop.length} noop`);
|
||||
|
||||
// Short-circuit: if every proposed file was identical to what's on
|
||||
// disk, there's nothing to test or PR. Clean up the branch and exit.
|
||||
if (changed.length === 0) {
|
||||
git(["checkout", "main"]);
|
||||
git(["branch", "-D", branch]);
|
||||
return finish("cycle_noop",
|
||||
`all ${proposal.files.length} proposed file(s) were identical to current state`,
|
||||
{ gap, proposal, filesNoop: ap.noop,
|
||||
tokens_used: proposal.tokens_used, cloud_calls: 1 });
|
||||
}
|
||||
|
||||
// 10. Test.
|
||||
log(`running tests (this will take a while)...`);
|
||||
const testRes = await runTests();
|
||||
if (!testRes.green) {
|
||||
git(["checkout", "."]);
|
||||
git(["checkout", "main"]);
|
||||
git(["branch", "-D", branch]);
|
||||
return finish("tests_failed", "cargo or bun test failed", {
|
||||
gap, proposal,
|
||||
filesAdded: ap.added, filesUpdated: ap.updated, filesNoop: ap.noop,
|
||||
testsGreen: false, testsOutput: testRes.output,
|
||||
tokens_used: proposal.tokens_used, cloud_calls: 1,
|
||||
});
|
||||
}
|
||||
log(`tests green`);
|
||||
|
||||
// 11. PR gate.
|
||||
const partial: CycleResult = {
|
||||
cycle_id, started_at,
|
||||
ended_at: new Date().toISOString(),
|
||||
outcome: "ok", reason: "",
|
||||
gap, proposal,
|
||||
filesAdded: ap.added, filesUpdated: ap.updated, filesNoop: ap.noop,
|
||||
testsGreen: true, testsOutput: testRes.output,
|
||||
prUrl: null, tokens_used: proposal.tokens_used, cloud_calls: 1,
|
||||
};
|
||||
const prGate = policy.shouldOpenPR(partial);
|
||||
if (!prGate.open) {
|
||||
git(["checkout", "main"]);
|
||||
git(["branch", "-D", branch]);
|
||||
return finish("pr_skipped_by_policy", prGate.reason, partial);
|
||||
}
|
||||
|
||||
// 12. Commit + push.
|
||||
git(["add", ...changed]);
|
||||
const commitMsg = [
|
||||
`bot: ${proposal.summary}`,
|
||||
"",
|
||||
proposal.rationale,
|
||||
"",
|
||||
`PRD gap: ${gap.prd_line}`,
|
||||
`Cycle: ${cycle_id}`,
|
||||
"",
|
||||
`Co-Authored-By: ${proposal.model_used} (cloud) <bot@lakehouse.local>`,
|
||||
].join("\n");
|
||||
const ci = git(["commit", "-m", commitMsg]);
|
||||
if (ci.code !== 0) {
|
||||
return finish("pr_failed", `git commit failed: ${ci.stderr}`, partial);
|
||||
}
|
||||
const ps = git(["push", "-u", "origin", branch]);
|
||||
if (ps.code !== 0) {
|
||||
return finish("pr_failed", `git push failed: ${ps.stderr}`, partial);
|
||||
}
|
||||
|
||||
// 13. Open PR.
|
||||
try {
|
||||
const pr = await openPr({
|
||||
branch,
|
||||
title: `bot: ${proposal.summary}`,
|
||||
body: [
|
||||
`**Gap** (PRD line ${gap.line_number}): ${gap.prd_line}`,
|
||||
"",
|
||||
`**Rationale**`,
|
||||
proposal.rationale,
|
||||
"",
|
||||
`**Cycle**: \`${cycle_id}\``,
|
||||
`**Model**: \`${proposal.model_used}\``,
|
||||
`**Tokens**: ${proposal.tokens_used}`,
|
||||
`**Added**: ${ap.added.map(f => `\`${f}\``).join(", ") || "_none_"}`,
|
||||
`**Updated**: ${ap.updated.map(f => `\`${f}\``).join(", ") || "_none_"}`,
|
||||
ap.noop.length > 0 ? `**NOOP (identical to current)**: ${ap.noop.map(f => `\`${f}\``).join(", ")}` : "",
|
||||
].join("\n"),
|
||||
});
|
||||
return finish("ok", `PR #${pr.number} opened`, { ...partial, prUrl: pr.html_url });
|
||||
} catch (e) {
|
||||
return finish("pr_failed", (e as Error).message, partial);
|
||||
}
|
||||
}
|
||||
|
||||
main().catch(e => {
|
||||
console.error("[bot] fatal:", e);
|
||||
process.exit(1);
|
||||
});
|
||||
@@ -1,101 +0,0 @@
|
||||
// Bot-local knowledge base. Every finished cycle already persists a
|
||||
// CycleResult to data/_bot/cycles/{id}.json — that IS the outcome log.
|
||||
// KB here just reads that dir, filters to prior cycles on the same gap,
|
||||
// and produces a short summary the cloud model can condition on.
|
||||
//
|
||||
// No separate jsonl, no new write path, no embedding calls. The bot's
|
||||
// "memory" is the same primary artifact that the observer consumes.
|
||||
//
|
||||
// Future: embedding-based neighbor matching across gaps (cheap once
|
||||
// sidecar is local), cross-pollination with scenario KB's
|
||||
// pathway_recommendations. Not required for the feedback loop to work
|
||||
// on a single gap — that's the floor we're building first.
|
||||
|
||||
import { readdir, readFile } from "node:fs/promises";
|
||||
import { join } from "node:path";
|
||||
import type { CycleResult } from "./types.ts";
|
||||
|
||||
const CYCLES_DIR = "/home/profit/lakehouse/data/_bot/cycles";
|
||||
|
||||
export interface HistoryEntry {
|
||||
cycle_id: string;
|
||||
ended_at: string;
|
||||
outcome: string;
|
||||
reason: string;
|
||||
pr_url: string | null;
|
||||
tests_green: boolean | null;
|
||||
files_added: string[];
|
||||
files_updated: string[];
|
||||
tokens_used: number;
|
||||
}
|
||||
|
||||
export async function loadHistory(gap_id: string, max: number = 5): Promise<HistoryEntry[]> {
|
||||
let entries: string[] = [];
|
||||
try {
|
||||
entries = await readdir(CYCLES_DIR);
|
||||
} catch {
|
||||
return [];
|
||||
}
|
||||
const matches: HistoryEntry[] = [];
|
||||
for (const e of entries) {
|
||||
if (!e.endsWith(".json")) continue;
|
||||
try {
|
||||
const raw = await readFile(join(CYCLES_DIR, e), "utf8");
|
||||
const r = JSON.parse(raw) as CycleResult;
|
||||
if (r.gap?.id !== gap_id) continue;
|
||||
matches.push({
|
||||
cycle_id: r.cycle_id,
|
||||
ended_at: r.ended_at,
|
||||
outcome: r.outcome,
|
||||
reason: r.reason,
|
||||
pr_url: r.prUrl,
|
||||
tests_green: r.testsGreen,
|
||||
files_added: r.filesAdded ?? [],
|
||||
files_updated: r.filesUpdated ?? [],
|
||||
tokens_used: r.tokens_used,
|
||||
});
|
||||
} catch {
|
||||
// Skip unreadable / malformed cycle files. Don't fail the current
|
||||
// cycle because an old one is corrupt.
|
||||
}
|
||||
}
|
||||
matches.sort((a, b) => b.ended_at.localeCompare(a.ended_at));
|
||||
return matches.slice(0, max);
|
||||
}
|
||||
|
||||
// Compact prompt-ready summary. Empty string when there's no history —
|
||||
// caller can skip the "prior attempts" block entirely.
|
||||
export function summarizeHistory(h: HistoryEntry[]): string {
|
||||
if (h.length === 0) return "";
|
||||
const lines = h.map(e => {
|
||||
const when = e.ended_at.slice(0, 16).replace("T", " ");
|
||||
const files = [...e.files_added, ...e.files_updated];
|
||||
const filesStr = files.length > 0 ? ` touched: ${files.join(", ")}` : "";
|
||||
const prStr = e.pr_url ? ` PR: ${e.pr_url}` : "";
|
||||
return `- ${when} UTC — ${e.outcome}${prStr}${filesStr}\n reason: ${e.reason}`;
|
||||
});
|
||||
return [
|
||||
`Prior attempts on this gap (${h.length} most recent):`,
|
||||
...lines,
|
||||
"",
|
||||
"Learn from these: build on what worked, avoid paths that failed.",
|
||||
].join("\n");
|
||||
}
|
||||
|
||||
// Aggregate stats for telemetry — lets the bot expose "% of cycles on
|
||||
// this gap that landed a PR" without re-parsing the raw history.
|
||||
export function statsFor(h: HistoryEntry[]): {
|
||||
attempts: number;
|
||||
pr_opened: number;
|
||||
tests_failed: number;
|
||||
proposal_rejected: number;
|
||||
noop: number;
|
||||
} {
|
||||
return {
|
||||
attempts: h.length,
|
||||
pr_opened: h.filter(e => e.outcome === "ok").length,
|
||||
tests_failed: h.filter(e => e.outcome === "tests_failed").length,
|
||||
proposal_rejected: h.filter(e => e.outcome === "proposal_rejected").length,
|
||||
noop: h.filter(e => e.outcome === "cycle_noop").length,
|
||||
};
|
||||
}
|
||||
@@ -1,34 +0,0 @@
|
||||
// POST bot cycle outcomes to the observer on :3800 so they accumulate
|
||||
// in the KB alongside scenario + autotune events. Non-fatal on failure:
|
||||
// the cycle result is also written to disk so observability isn't a
|
||||
// single point of failure.
|
||||
|
||||
const OBSERVER_URL = process.env.LH_OBSERVER_URL ?? "http://localhost:3800";
|
||||
|
||||
export interface BotObserverEvent {
|
||||
source: "bot";
|
||||
cycle_id: string;
|
||||
sig_hash: string; // stable hash of the gap + proposal for dedup
|
||||
event_kind: string; // cycle outcome (ok, tests_failed, ...)
|
||||
ok: boolean;
|
||||
staffer_id?: string; // the model used, e.g. "gpt-oss:120b"
|
||||
turns?: number; // number of cloud calls this cycle
|
||||
duration_ms?: number;
|
||||
extra?: Record<string, any>;
|
||||
}
|
||||
|
||||
export async function postEvent(ev: BotObserverEvent): Promise<void> {
|
||||
try {
|
||||
const r = await fetch(`${OBSERVER_URL}/event`, {
|
||||
method: "POST",
|
||||
headers: { "content-type": "application/json" },
|
||||
body: JSON.stringify(ev),
|
||||
signal: AbortSignal.timeout(3000),
|
||||
});
|
||||
if (!r.ok) {
|
||||
console.error(`[bot/observer] ${r.status}: ${await r.text()}`);
|
||||
}
|
||||
} catch (e) {
|
||||
console.error(`[bot/observer] POST failed: ${(e as Error).message}`);
|
||||
}
|
||||
}
|
||||
@@ -1,70 +0,0 @@
|
||||
// ═══════════════════════════════════════════════════════════════════
|
||||
// YOU WRITE THIS FILE. These four functions shape the bot's entire
|
||||
// behavior. Defaults below are safe and intentionally conservative —
|
||||
// redline to taste. Each function is ~5-10 lines.
|
||||
// ═══════════════════════════════════════════════════════════════════
|
||||
|
||||
import type { CycleContext, CycleResult, Gap, Proposal } from "./types.ts";
|
||||
|
||||
// Should this cycle even fire right now?
|
||||
// Called AFTER the pause-file check + cost-cap check have already passed.
|
||||
// Use this for judgment calls (cooldown, shared-resource conflicts).
|
||||
export function shouldRunCycle(ctx: CycleContext): { run: boolean; reason: string } {
|
||||
if (ctx.workingTreeDirty) {
|
||||
return { run: false, reason: "dirty working tree — commit or stash your work first" };
|
||||
}
|
||||
if (ctx.lastCycleAt) {
|
||||
const minsAgo = (Date.now() - new Date(ctx.lastCycleAt).getTime()) / 60000;
|
||||
if (minsAgo < 10) return { run: false, reason: `last cycle ${minsAgo.toFixed(1)}min ago (<10min cooldown)` };
|
||||
}
|
||||
if (ctx.autotuneBusy) {
|
||||
return { run: false, reason: "autotune agent is running — avoid Ollama contention" };
|
||||
}
|
||||
return { run: true, reason: "ready" };
|
||||
}
|
||||
|
||||
// Which [bot-eligible] gap from the PRD does this cycle attack?
|
||||
// Default: first in document order. Alternatives: random for exploration,
|
||||
// smallest-context, tag-priority (e.g. [bot-eligible:high] wins).
|
||||
export function pickGap(gaps: Gap[]): { gap: Gap | null; reason: string } {
|
||||
if (gaps.length === 0) return { gap: null, reason: "no [bot-eligible] gaps in PRD" };
|
||||
return { gap: gaps[0], reason: "first eligible gap in document order" };
|
||||
}
|
||||
|
||||
// Cheap gate BEFORE burning test cycles. Reject if proposal smells wrong.
|
||||
// This catches runaway cloud proposals before we apply+test them.
|
||||
export function scoreProposal(p: Proposal): { accept: boolean; reason: string } {
|
||||
if (p.files.length === 0) return { accept: false, reason: "empty proposal" };
|
||||
if (p.files.length > 5) return { accept: false, reason: `>5 files (${p.files.length})` };
|
||||
if (p.estimated_loc > 200) return { accept: false, reason: `>200 LOC (${p.estimated_loc})` };
|
||||
|
||||
// Paths the bot is NEVER allowed to touch. Extend cautiously.
|
||||
const forbidden = [
|
||||
".git/", "secrets", "lakehouse.toml",
|
||||
"docs/ADR-", "docs/DECISIONS.md", "docs/PRD.md",
|
||||
"/etc/", "/root/", "Cargo.lock",
|
||||
];
|
||||
for (const f of p.files) {
|
||||
const hit = forbidden.find(prefix => f.path.includes(prefix));
|
||||
if (hit) return { accept: false, reason: `touches forbidden path (${hit}): ${f.path}` };
|
||||
}
|
||||
return { accept: true, reason: "passes size + path filters" };
|
||||
}
|
||||
|
||||
// Tests are green and changes applied. Do we actually open a PR?
|
||||
// Mem0-aligned: the cycle already short-circuited with `cycle_noop` if
|
||||
// every file was identical to current state, so by the time we're here
|
||||
// at least one file was actually added or updated.
|
||||
//
|
||||
// Stricter options below — uncomment to require a test file touched, etc.
|
||||
export function shouldOpenPR(c: CycleResult): { open: boolean; reason: string } {
|
||||
if (!c.testsGreen) return { open: false, reason: "tests red" };
|
||||
const realChanges = c.filesAdded.length + c.filesUpdated.length;
|
||||
if (realChanges === 0) return { open: false, reason: "nothing actually changed on disk" };
|
||||
|
||||
// Stricter option — require at least one test file in the diff:
|
||||
// const hasTest = [...c.filesAdded, ...c.filesUpdated].some(f => /test|_tests/i.test(f));
|
||||
// if (!hasTest) return { open: false, reason: "no test file touched" };
|
||||
|
||||
return { open: true, reason: `green + ${c.filesAdded.length} add / ${c.filesUpdated.length} update` };
|
||||
}
|
||||
@@ -1,58 +0,0 @@
|
||||
// Gitea PR open. Reads PAT from ~/.git-credentials (set up by the
|
||||
// credential-helper flow). All PRs open as DRAFT so J reviews before
|
||||
// merge — never auto-merge.
|
||||
|
||||
import { readFile } from "node:fs/promises";
|
||||
|
||||
const GITEA_HOST = "https://git.agentview.dev";
|
||||
const REPO_OWNER = "profit";
|
||||
const REPO_NAME = "lakehouse";
|
||||
const CRED_FILE = "/home/profit/.git-credentials";
|
||||
|
||||
async function readPat(): Promise<string> {
|
||||
const raw = await readFile(CRED_FILE, "utf8");
|
||||
for (const line of raw.split("\n")) {
|
||||
const m = line.match(/^https:\/\/[^:]+:([^@]+)@git\.agentview\.dev/);
|
||||
if (m) return m[1];
|
||||
}
|
||||
throw new Error(`no Gitea PAT found in ${CRED_FILE}`);
|
||||
}
|
||||
|
||||
export interface OpenPrInput {
|
||||
branch: string; // head branch (just pushed)
|
||||
base?: string; // default "main"
|
||||
title: string;
|
||||
body: string;
|
||||
}
|
||||
|
||||
export interface OpenPrResult {
|
||||
number: number;
|
||||
html_url: string;
|
||||
}
|
||||
|
||||
export async function openPr(input: OpenPrInput): Promise<OpenPrResult> {
|
||||
const pat = await readPat();
|
||||
const url = `${GITEA_HOST}/api/v1/repos/${REPO_OWNER}/${REPO_NAME}/pulls`;
|
||||
const payload = {
|
||||
title: input.title,
|
||||
body: input.body,
|
||||
head: input.branch,
|
||||
base: input.base ?? "main",
|
||||
// Gitea uses "draft" flag — always on for bot PRs.
|
||||
draft: true,
|
||||
};
|
||||
const r = await fetch(url, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
"content-type": "application/json",
|
||||
"Authorization": `token ${pat}`,
|
||||
},
|
||||
body: JSON.stringify(payload),
|
||||
signal: AbortSignal.timeout(30000),
|
||||
});
|
||||
if (!r.ok) {
|
||||
throw new Error(`Gitea ${r.status}: ${await r.text()}`);
|
||||
}
|
||||
const j = await r.json() as any;
|
||||
return { number: j.number, html_url: j.html_url };
|
||||
}
|
||||
-141
@@ -1,141 +0,0 @@
|
||||
// Gap detection + cloud proposal.
|
||||
//
|
||||
// Gap detection: scan docs/PRD.md for lines tagged [bot-eligible].
|
||||
// Each match becomes a Gap with surrounding context.
|
||||
//
|
||||
// Proposal: one-shot call to the T3 cloud model via the Python sidecar's
|
||||
// /generate endpoint. Asks for a structured JSON response with file
|
||||
// contents. Truncation-resistant via Phase 21's generate_continuable —
|
||||
// for now we pass max_tokens high and rely on the model completing in
|
||||
// one pass; swap to the Rust continuation wrapper if we see truncation.
|
||||
|
||||
import { readFile } from "node:fs/promises";
|
||||
import { createHash } from "node:crypto";
|
||||
import type { Gap, Proposal } from "./types.ts";
|
||||
|
||||
const SIDECAR_URL = process.env.LH_SIDECAR_URL ?? "http://localhost:3200";
|
||||
const REPO_ROOT = "/home/profit/lakehouse";
|
||||
const PRD_PATH = `${REPO_ROOT}/docs/PRD.md`;
|
||||
const CLOUD_MODEL = process.env.LH_BOT_MODEL ?? "gpt-oss:120b";
|
||||
const MAX_TOKENS = 6000;
|
||||
|
||||
export async function findGaps(): Promise<Gap[]> {
|
||||
const prd = await readFile(PRD_PATH, "utf8");
|
||||
const lines = prd.split("\n");
|
||||
const gaps: Gap[] = [];
|
||||
for (let i = 0; i < lines.length; i++) {
|
||||
if (!lines[i].includes("[bot-eligible]")) continue;
|
||||
const contextLines = lines.slice(i, Math.min(i + 6, lines.length)).join("\n");
|
||||
const id = createHash("sha256").update(lines[i]).digest("hex").slice(0, 12);
|
||||
gaps.push({
|
||||
id,
|
||||
prd_line: lines[i].trim(),
|
||||
context: contextLines,
|
||||
source_file: "docs/PRD.md",
|
||||
line_number: i + 1,
|
||||
});
|
||||
}
|
||||
return gaps;
|
||||
}
|
||||
|
||||
const SYSTEM_PROMPT = `You are an assistant that proposes small, testable code changes to the Lakehouse repo.
|
||||
The Lakehouse is a Rust-first data platform with 13 crates + Bun/TypeScript test harness.
|
||||
You will be given one PRD gap tagged [bot-eligible] and must respond with a STRICT JSON object — no prose.
|
||||
|
||||
Rules:
|
||||
- Response MUST be a single JSON object, no markdown fences, no commentary.
|
||||
- Change MUST be small: <200 lines total, ≤5 files.
|
||||
- Include at least one test file (new or modified) that proves the change.
|
||||
- NEVER touch .git/, secrets, lakehouse.toml, docs/ADR-*, docs/DECISIONS.md, docs/PRD.md, /etc/, /root/, Cargo.lock.
|
||||
- Paths MUST be repo-relative (no leading /).
|
||||
- Whole-file contents only — no patches, no diffs.
|
||||
|
||||
Response shape:
|
||||
{
|
||||
"summary": "one line",
|
||||
"rationale": "why this addresses the gap",
|
||||
"files": [ { "path": "crates/foo/src/bar.rs", "content": "<full file>", "is_new": false } ],
|
||||
"estimated_loc": 42
|
||||
}`;
|
||||
|
||||
export async function generateProposal(gap: Gap, historySummary: string = ""): Promise<Proposal> {
|
||||
const sections = [
|
||||
`PRD gap (line ${gap.line_number}):`,
|
||||
"```",
|
||||
gap.context,
|
||||
"```",
|
||||
"",
|
||||
];
|
||||
if (historySummary) {
|
||||
sections.push(historySummary, "");
|
||||
}
|
||||
sections.push("Propose a small change that addresses this gap. Respond with the JSON object only.");
|
||||
const userPrompt = sections.join("\n");
|
||||
|
||||
const r = await fetch(`${SIDECAR_URL}/generate`, {
|
||||
method: "POST",
|
||||
headers: { "content-type": "application/json" },
|
||||
body: JSON.stringify({
|
||||
model: CLOUD_MODEL,
|
||||
system: SYSTEM_PROMPT,
|
||||
prompt: userPrompt,
|
||||
temperature: 0.2,
|
||||
max_tokens: MAX_TOKENS,
|
||||
think: false,
|
||||
}),
|
||||
signal: AbortSignal.timeout(180000), // cloud T3 can be slow — 3 min
|
||||
});
|
||||
if (!r.ok) {
|
||||
throw new Error(`sidecar ${r.status}: ${await r.text()}`);
|
||||
}
|
||||
const j = await r.json() as any;
|
||||
const raw: string = j.text ?? j.response ?? "";
|
||||
const usage = j.usage ?? {};
|
||||
const tokens = (usage.prompt_tokens ?? 0) + (usage.completion_tokens ?? 0);
|
||||
|
||||
const parsed = extractJson(raw);
|
||||
if (!parsed || typeof parsed !== "object") {
|
||||
throw new Error(`model returned no JSON object. Raw head: ${raw.slice(0, 300)}`);
|
||||
}
|
||||
if (!Array.isArray(parsed.files)) {
|
||||
throw new Error(`proposal.files not an array: ${JSON.stringify(parsed).slice(0, 200)}`);
|
||||
}
|
||||
|
||||
return {
|
||||
summary: String(parsed.summary ?? "").trim(),
|
||||
rationale: String(parsed.rationale ?? "").trim(),
|
||||
files: parsed.files.map((f: any) => ({
|
||||
path: String(f.path ?? ""),
|
||||
content: String(f.content ?? ""),
|
||||
is_new: Boolean(f.is_new),
|
||||
})).filter((f: any) => f.path && f.content !== undefined),
|
||||
estimated_loc: Number(parsed.estimated_loc ?? 0),
|
||||
model_used: CLOUD_MODEL,
|
||||
tokens_used: tokens,
|
||||
};
|
||||
}
|
||||
|
||||
// Find the first balanced JSON object in the string. Tolerates leading
|
||||
// "```json" fences even though we asked the model not to emit them.
|
||||
function extractJson(text: string): any | null {
|
||||
const cleaned = text.replace(/^```(?:json)?\s*/i, "").replace(/```\s*$/i, "").trim();
|
||||
let depth = 0;
|
||||
let start = -1;
|
||||
for (let i = 0; i < cleaned.length; i++) {
|
||||
const c = cleaned[i];
|
||||
if (c === "{") {
|
||||
if (depth === 0) start = i;
|
||||
depth++;
|
||||
} else if (c === "}") {
|
||||
depth--;
|
||||
if (depth === 0 && start >= 0) {
|
||||
try {
|
||||
return JSON.parse(cleaned.slice(start, i + 1));
|
||||
} catch {
|
||||
start = -1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
-46
@@ -1,46 +0,0 @@
|
||||
// Run the test gate: `cargo test --workspace` + `bun test`.
|
||||
// Returns { green, output } — output is last ~4KB of combined stdout/stderr.
|
||||
|
||||
import { spawn } from "node:child_process";
|
||||
|
||||
const REPO_ROOT = "/home/profit/lakehouse";
|
||||
const TEST_TIMEOUT_MS = 15 * 60 * 1000; // 15 min cargo + 2 min bun, generous
|
||||
|
||||
export async function runTests(): Promise<{ green: boolean; output: string }> {
|
||||
const cargoOut = await runCmd("cargo", ["test", "--workspace", "--quiet", "--", "--test-threads=1"], REPO_ROOT, TEST_TIMEOUT_MS);
|
||||
if (cargoOut.code !== 0) {
|
||||
return { green: false, output: tail(`cargo test (${cargoOut.code}):\n${cargoOut.combined}`) };
|
||||
}
|
||||
|
||||
const bunOut = await runCmd("bun", ["test", "tests/multi-agent"], REPO_ROOT, 120000);
|
||||
const bunGreen = bunOut.code === 0;
|
||||
const combined = [
|
||||
`cargo test: OK`,
|
||||
`bun test (${bunOut.code}):`,
|
||||
bunOut.combined,
|
||||
].join("\n");
|
||||
return { green: bunGreen, output: tail(combined) };
|
||||
}
|
||||
|
||||
function tail(s: string, n = 4096): string {
|
||||
return s.length > n ? "…" + s.slice(-n) : s;
|
||||
}
|
||||
|
||||
function runCmd(
|
||||
cmd: string,
|
||||
args: string[],
|
||||
cwd: string,
|
||||
timeoutMs: number,
|
||||
): Promise<{ code: number; combined: string }> {
|
||||
return new Promise(resolveP => {
|
||||
const child = spawn(cmd, args, { cwd, env: { ...process.env } });
|
||||
let combined = "";
|
||||
child.stdout.on("data", d => { combined += d.toString(); });
|
||||
child.stderr.on("data", d => { combined += d.toString(); });
|
||||
const timer = setTimeout(() => child.kill("SIGKILL"), timeoutMs);
|
||||
child.on("close", code => {
|
||||
clearTimeout(timer);
|
||||
resolveP({ code: code ?? -1, combined });
|
||||
});
|
||||
});
|
||||
}
|
||||
@@ -1,88 +0,0 @@
|
||||
// Shared types for the PR-bot. Small and explicit — the bot's behavior
|
||||
// should be readable from these shapes alone.
|
||||
|
||||
export interface Gap {
|
||||
id: string; // hash of prd_line — stable key for dedup
|
||||
prd_line: string; // the [bot-eligible]-tagged line, verbatim
|
||||
context: string; // surrounding context from the PRD (next ~5 lines)
|
||||
source_file: string; // e.g. "docs/PRD.md"
|
||||
line_number: number;
|
||||
}
|
||||
|
||||
export interface ProposalFile {
|
||||
path: string; // repo-relative (no leading slash)
|
||||
content: string; // full file content — the bot writes whole files, not patches
|
||||
is_new: boolean; // true = create, false = overwrite existing
|
||||
}
|
||||
|
||||
export interface Proposal {
|
||||
summary: string; // one-line
|
||||
rationale: string; // why this change addresses the gap
|
||||
files: ProposalFile[];
|
||||
estimated_loc: number; // total added+changed lines across all files
|
||||
model_used: string;
|
||||
tokens_used: number;
|
||||
}
|
||||
|
||||
export interface CycleContext {
|
||||
startedAt: string; // ISO
|
||||
dailyCallsUsed: number;
|
||||
dailyCallsBudget: number;
|
||||
dailyTokensUsed: number;
|
||||
dailyTokensBudget: number;
|
||||
lastCycleAt: string | null;
|
||||
lastCycleGapId: string | null;
|
||||
autotuneBusy: boolean;
|
||||
workingTreeDirty: boolean;
|
||||
}
|
||||
|
||||
export type CycleOutcome =
|
||||
| "ok" // PR opened
|
||||
| "skipped_pause"
|
||||
| "skipped_cost"
|
||||
| "skipped_policy"
|
||||
| "skipped_no_gap"
|
||||
| "skipped_dirty_tree"
|
||||
| "proposal_rejected"
|
||||
| "apply_failed"
|
||||
| "cycle_noop" // proposal applied but every file was identical to what's on disk
|
||||
| "tests_failed"
|
||||
| "pr_skipped_by_policy"
|
||||
| "pr_failed"
|
||||
| "model_failed";
|
||||
|
||||
// Mem0-aligned apply outcomes. Three shapes instead of the binary
|
||||
// "written" / "errored". NOOP is a first-class outcome — identical
|
||||
// content shouldn't waste test cycles or open an empty PR.
|
||||
export type ApplyMode = "add" | "update" | "noop";
|
||||
|
||||
export interface ApplyOutcome {
|
||||
added: string[];
|
||||
updated: string[];
|
||||
noop: string[];
|
||||
errors: string[];
|
||||
}
|
||||
|
||||
export interface CycleResult {
|
||||
cycle_id: string;
|
||||
started_at: string;
|
||||
ended_at: string;
|
||||
outcome: CycleOutcome;
|
||||
reason: string;
|
||||
gap: Gap | null;
|
||||
proposal: Proposal | null;
|
||||
filesAdded: string[];
|
||||
filesUpdated: string[];
|
||||
filesNoop: string[];
|
||||
testsGreen: boolean | null; // null = not run
|
||||
testsOutput: string;
|
||||
prUrl: string | null;
|
||||
tokens_used: number;
|
||||
cloud_calls: number;
|
||||
}
|
||||
|
||||
export interface CostState {
|
||||
date: string; // YYYY-MM-DD UTC
|
||||
calls: number;
|
||||
tokens: number;
|
||||
}
|
||||
@@ -1,18 +0,0 @@
|
||||
{
|
||||
"lockfileVersion": 1,
|
||||
"configVersion": 1,
|
||||
"workspaces": {
|
||||
"": {
|
||||
"dependencies": {
|
||||
"langfuse": "^3.38.20",
|
||||
},
|
||||
},
|
||||
},
|
||||
"packages": {
|
||||
"langfuse": ["langfuse@3.38.20", "", { "dependencies": { "langfuse-core": "^3.38.20" } }, "sha512-MAmBAASSzJtmK1O9HQegA1mFsQhT8Yf+OJRGvE7FXkyv3g/eiBE0glLD0Ohg3pkxhoPdggM5SejK7ue9ctlaMA=="],
|
||||
|
||||
"langfuse-core": ["langfuse-core@3.38.20", "", { "dependencies": { "mustache": "^4.2.0" } }, "sha512-zBKVmQN/1oT5VWZUBYlWzvokIlkC/6mnpgr/2atMyTeAm+jR3ia7w2iJMjlrF5/oG8ukO1s8+LDRCzJpF1QeEA=="],
|
||||
|
||||
"mustache": ["mustache@4.2.0", "", { "bin": { "mustache": "bin/mustache" } }, "sha512-71ippSywq5Yb7/tVYyGbkBggbU8H3u5Rz56fH60jGFgr8uHwxs+aSKeqmluIVzM0m0kB7xQjKS6qPfd0b2ZoqQ=="],
|
||||
}
|
||||
}
|
||||
@@ -1,193 +0,0 @@
|
||||
{
|
||||
"_description": "Lakehouse model matrix — authoritative routing for all agent tiers. Local models do the heavy lifting; cloud models are consulted sparingly for overview, strategic, and gatekeeper decisions. Read by tests/multi-agent/scenario.ts and mcp-server/index.ts.",
|
||||
"version": 1,
|
||||
"updated": "2026-04-21",
|
||||
"providers": {
|
||||
"ollama_local": {
|
||||
"base_url": "http://localhost:11434",
|
||||
"key_env": null
|
||||
},
|
||||
"ollama_cloud": {
|
||||
"base_url": "https://ollama.com",
|
||||
"key_env": "OLLAMA_CLOUD_KEY",
|
||||
"key_source": "/root/llm_team_config.json → providers.ollama_cloud.api_key",
|
||||
"rate_budget": {
|
||||
"calls_per_hour": 200,
|
||||
"calls_per_day": 2000,
|
||||
"notes": "Paid tier — generous. Policy: keep overview calls ≤ 3/scenario, strategic ≤ 10/day, gatekeeper ≤ 5/day."
|
||||
}
|
||||
}
|
||||
},
|
||||
"tiers": {
|
||||
"t1_hot": {
|
||||
"purpose": "Per tool call — SQL generation, hybrid_search, sql(). Runs 50-200 times per scenario. Latency-sensitive.",
|
||||
"kind": "local_fast",
|
||||
"primary": { "model": "qwen3.5:latest", "provider": "ollama_local", "context_window": 262144 },
|
||||
"fallback": { "model": "qwen2.5:latest", "provider": "ollama_local", "context_window": 32768 },
|
||||
"max_tokens": 1000,
|
||||
"temperature": 0.3,
|
||||
"never_route_cloud": true,
|
||||
"context_budget": {
|
||||
"system_prompt_cap": 4000,
|
||||
"prior_context_cap": 6000,
|
||||
"tool_results_cap": 8000,
|
||||
"safety_margin": 2000,
|
||||
"overflow_policy": "summarize_oldest_tool_results_via_t3"
|
||||
},
|
||||
"rationale": "qwen3.5:latest is a 9.7B thinking model with 262K context that emits clean JSON (verified 2026-04-21 after mistral A/B wipeout). Thinking budget requires max_tokens ≥ 400; we hold 1000 to keep propose_done payloads intact. qwen2.5 stays as fallback."
|
||||
},
|
||||
"t2_review": {
|
||||
"purpose": "Per step consensus — executor ↔ reviewer loop critique. 5-14 calls per event.",
|
||||
"kind": "local_balanced",
|
||||
"primary": { "model": "qwen3:latest", "provider": "ollama_local", "context_window": 40960 },
|
||||
"fallback": { "model": "qwen2.5:latest", "provider": "ollama_local", "context_window": 32768 },
|
||||
"max_tokens": 800,
|
||||
"temperature": 0.3,
|
||||
"never_route_cloud": true,
|
||||
"context_budget": {
|
||||
"system_prompt_cap": 2000,
|
||||
"recent_turns_cap": 4000,
|
||||
"safety_margin": 1000
|
||||
},
|
||||
"rationale": "Reviewer only needs to detect schema violations and drift — a 7B model is sufficient."
|
||||
},
|
||||
"t3_overview": {
|
||||
"purpose": "Mid-day checkpoint after every misplacement + every Nth event, and cross-day lesson. 1-3 calls per scenario.",
|
||||
"kind": "thinking_cloud",
|
||||
"primary": { "model": "gpt-oss:120b", "provider": "ollama_cloud", "context_window": 131072 },
|
||||
"local_fallback": { "model": "gpt-oss:20b", "provider": "ollama_local", "context_window": 131072 },
|
||||
"max_tokens": 900,
|
||||
"temperature": 0.2,
|
||||
"cloud_budget_per_scenario": 5,
|
||||
"env_flag": "LH_OVERVIEW_CLOUD=1",
|
||||
"context_budget": {
|
||||
"event_digest_cap": 30000,
|
||||
"checkpoint_cap": 8000,
|
||||
"lesson_corpus_cap": 40000,
|
||||
"safety_margin": 8000,
|
||||
"overflow_policy": "chunk_lessons_via_cosine_topk"
|
||||
},
|
||||
"rationale": "Same prompt family as local 20b (gpt-oss series) — prompts port directly. 120b is faster via cloud than 20b local in practice, and lessons are noticeably more specific."
|
||||
},
|
||||
"t4_strategic": {
|
||||
"purpose": "Daily playbook board re-ranking, weekly gap audit, pattern discovery across accumulated playbooks. 1-10 calls per day.",
|
||||
"kind": "thinking_cloud_large",
|
||||
"primary": { "model": "kimi-k2.6", "provider": "ollama_cloud", "context_window": 200000 },
|
||||
"secondary": { "model": "qwen3.5:397b", "provider": "ollama_cloud", "context_window": 131072 },
|
||||
"fallback": { "model": "glm-4.7", "provider": "ollama_cloud", "context_window": 131072 },
|
||||
"local_fallback": { "model": "gpt-oss:20b", "provider": "ollama_local", "context_window": 131072 },
|
||||
"max_tokens": 2000,
|
||||
"temperature": 0.2,
|
||||
"cloud_budget_per_day": 10,
|
||||
"context_budget": {
|
||||
"playbook_corpus_cap": 80000,
|
||||
"pattern_history_cap": 20000,
|
||||
"safety_margin": 16000,
|
||||
"overflow_policy": "two_pass_map_reduce"
|
||||
},
|
||||
"rationale": "J named qwen3.5 specifically. GLM-4.7 is a promising alternate for debate phase. Runs after all scenarios complete for the day."
|
||||
},
|
||||
"t5_gatekeeper": {
|
||||
"purpose": "MUST route here: architecture changes, new client onboarding, schema migrations, playbook retirements, index rebuilds, autotune config changes.",
|
||||
"kind": "thinking_cloud_deepest",
|
||||
"primary": { "model": "kimi-k2-thinking", "provider": "ollama_cloud", "context_window": 200000 },
|
||||
"fallback": { "model": "deepseek-v3.1:671b", "provider": "ollama_cloud", "context_window": 131072 },
|
||||
"secondary_fallback": { "model": "qwen3.5:397b", "provider": "ollama_cloud", "context_window": 131072 },
|
||||
"local_fallback": { "model": "gpt-oss:20b", "provider": "ollama_local", "context_window": 131072 },
|
||||
"max_tokens": 4000,
|
||||
"temperature": 0.1,
|
||||
"cloud_budget_per_day": 5,
|
||||
"audit_log": true,
|
||||
"context_budget": {
|
||||
"decision_doc_cap": 50000,
|
||||
"evidence_bundle_cap": 100000,
|
||||
"prior_gatekeeper_decisions_cap": 20000,
|
||||
"safety_margin": 20000,
|
||||
"overflow_policy": "escalate_to_kimi_k2_1t_or_split_decision"
|
||||
},
|
||||
"rationale": "Highest-stakes decisions — reasoning depth matters more than latency. Audit log so J can always see what the gatekeeper was asked and what it answered. No human approval required today; escalate later if mis-decisions show up."
|
||||
}
|
||||
},
|
||||
|
||||
"embeddings": {
|
||||
"_description": "Vector embedding models for hybrid search + playbook_memory similarity. Always local — embedding calls are too high-volume for cloud. qwen3-embedding is the new primary (2026-04-21); nomic-embed-text-v2-moe is the MoE alternative for corpus-scale embedding; original nomic-embed-text stays as legacy fallback while vectord indexes are being re-embedded.",
|
||||
"primary": { "model": "qwen3-embedding", "provider": "ollama_local", "dim": 1024, "context_window": 32768 },
|
||||
"alternate": { "model": "nomic-embed-text-v2-moe", "provider": "ollama_local", "dim": 768, "context_window": 2048, "notes": "MoE variant — better for long documents, slower per-call, roughly matches nomic-embed-text on short text." },
|
||||
"legacy": { "model": "nomic-embed-text:latest", "provider": "ollama_local", "dim": 768, "context_window": 8192, "notes": "Used by existing vectord indexes. Do NOT switch live until indexes are rebuilt, or hybrid search recall will crater." },
|
||||
"migration_plan": "Re-embed workers_500k (and all production indexes) with qwen3-embedding, keep the old index under a _v1 suffix for rollback, flip the profile pointer only after recall benchmarks match or exceed legacy. Touches vectord::agent + autotune."
|
||||
},
|
||||
|
||||
"context_management": {
|
||||
"_description": "Rule zero: NEVER call a model with more tokens than its context_window minus safety_margin. Every call goes through the budget checker first. If over budget → chunk, summarize, or escalate. This is the stability floor.",
|
||||
"token_estimator": {
|
||||
"method": "chars_div_4",
|
||||
"note": "Rough, biased safe by ~15%. For production, swap to tiktoken or the provider's tokenizer endpoint."
|
||||
},
|
||||
"overflow_policies": {
|
||||
"summarize_oldest_tool_results_via_t3": {
|
||||
"when": "T1 conversation history + tool results exceed context_budget.tool_results_cap",
|
||||
"how": "Send oldest N tool results to T3 with prompt 'summarize these into 500 tokens that preserve what the executor needs to know'; replace them with the summary in the running conversation."
|
||||
},
|
||||
"chunk_lessons_via_cosine_topk": {
|
||||
"when": "lesson corpus in data/_playbook_lessons/*.json exceeds lesson_corpus_cap",
|
||||
"how": "Embed the current scenario spec, cosine-rank all lessons, take top-K until budget exhausted. Fall back to date-sorted if embeddings unavailable."
|
||||
},
|
||||
"two_pass_map_reduce": {
|
||||
"when": "T4 playbook corpus exceeds playbook_corpus_cap",
|
||||
"how": "Pass 1: chunk playbooks into ≤30K token shards, run primary model on each shard to emit 'shard summary'. Pass 2: feed all summaries to primary model for global synthesis. Logged as two audit entries."
|
||||
},
|
||||
"escalate_to_kimi_k2_1t_or_split_decision": {
|
||||
"when": "T5 decision evidence exceeds decision_doc_cap + evidence_bundle_cap",
|
||||
"how": "Prefer kimi-k2:1t which has 1M context. If still over, split decision into sub-decisions (e.g. 'retire playbooks by city' instead of 'retire playbooks globally') and loop."
|
||||
}
|
||||
},
|
||||
"chunking_cache": {
|
||||
"_description": "Precomputed shards of the playbook corpus, indexed by (corpus_version, shard_id). Avoids re-chunking on every T4 run.",
|
||||
"location": "data/_chunk_cache/",
|
||||
"invalidation": "Key includes corpus hash. When playbook_memory changes, the hash changes, the cache misses, and chunks regenerate.",
|
||||
"implementation_status": "SPEC — next sprint."
|
||||
},
|
||||
"implementation_status": "context_window + context_budget fields WIRED in config. Enforcement helper NOT yet wired in agent.ts. Next: add estimateTokens() + budgetCheck() helpers, route all generate() calls through them, emit a metric when overflow policy fires."
|
||||
},
|
||||
"experimental_rotation": {
|
||||
"enabled": false,
|
||||
"purpose": "Sample newer models on a schedule to collect comparison data without rate-limit risk.",
|
||||
"candidates": [
|
||||
{ "model": "minimax-m2.7", "notes": "Newer minimax; unknown output stability" },
|
||||
{ "model": "glm-5", "notes": "GLM next-gen; larger context" },
|
||||
{ "model": "glm-5.1", "notes": "Incremental on GLM-5" },
|
||||
{ "model": "qwen3-next:80b", "notes": "Qwen's experimental successor; smaller than 3.5" },
|
||||
{ "model": "qwen3-coder-next", "notes": "Coder-optimized — good for SQL gen T1 experiments" },
|
||||
{ "model": "deepseek-v3.2", "notes": "Smaller deepseek; reasoning/coding" },
|
||||
{ "model": "nemotron-3-super", "notes": "NVIDIA 230B; general" },
|
||||
{ "model": "cogito-2.1:671b", "notes": "671B general" },
|
||||
{ "model": "mistral-large-3:675b", "notes": "Mistral's flagship; good T3 candidate" }
|
||||
],
|
||||
"rotation": "weekly",
|
||||
"sample_rate": 0.1,
|
||||
"apply_to_tier": "t4_strategic",
|
||||
"notes": "When enabled, T4 routes 10% of calls to a rotating experimental model. Log comparison in /data/_model_eval/ — if the experimental consistently beats primary across 3 rotations, promote it to primary."
|
||||
},
|
||||
"playbook_versioning": {
|
||||
"enabled": true,
|
||||
"purpose": "A playbook can work, then break when architecture changes. Versioning lets us pin which change retired which playbook.",
|
||||
"dataset": "playbook_memory",
|
||||
"schema_additions": {
|
||||
"version": "integer — auto-increment per operation",
|
||||
"parent_id": "string — previous version entry_id for same operation (null for v1)",
|
||||
"retired_at": "timestamp — set when success_rate drops or architecture changes",
|
||||
"retirement_reason": "string — e.g. 'schema_migration:workers_500k 2026-05-03'",
|
||||
"architecture_snapshot": "object — crate versions, index name, schema fingerprint at seed time"
|
||||
},
|
||||
"retire_triggers": [
|
||||
"success_rate < 0.3 over last 20 citations",
|
||||
"schema_fingerprint mismatch detected at retrieval time",
|
||||
"architecture change event emitted by ingestd/vectord",
|
||||
"T5 gatekeeper explicitly retires via /vectors/playbook_memory/retire"
|
||||
],
|
||||
"read_back_policy": "retrieval returns only non-retired versions. History endpoint /vectors/playbook_memory/history/{operation} returns the full chain.",
|
||||
"ui_surface": "mcp-server to render a diff view: side-by-side of playbook versions with a timeline of what changed and when.",
|
||||
"implementation_status": "SPEC — not yet wired. Target: next sprint. Touches gateway + catalogd + mcp-server."
|
||||
},
|
||||
"matrix_index_hybrid_search_note": "Phase 22 candidate: elevate the hybrid_search T1 tool to consult T3 when a pool returns <3 matches OR when the same (role, city) combo has failed N times in 24h. Consult result is a reformulated sql_filter the executor retries with. Keeps T1 fast on the happy path, escalates to T3 only on low-recall signals."
|
||||
}
|
||||
@@ -1,61 +0,0 @@
|
||||
# Phase 40: Routing Engine Configuration
|
||||
#
|
||||
# Human-editable rules for model → provider routing.
|
||||
# Matching order: first match wins.
|
||||
|
||||
[[rule]]
|
||||
model_pattern = "gpt-4*"
|
||||
provider = "openrouter"
|
||||
max_tokens = 4096
|
||||
temperature = 0.7
|
||||
|
||||
[[rule]]
|
||||
model_pattern = "claude*"
|
||||
provider = "claude"
|
||||
max_tokens = 4096
|
||||
temperature = 0.7
|
||||
|
||||
[[rule]]
|
||||
model_pattern = "gemini*"
|
||||
provider = "gemini"
|
||||
max_tokens = 8192
|
||||
temperature = 0.9
|
||||
|
||||
[[rule]]
|
||||
model_pattern = "qwen3.5*"
|
||||
provider = "ollama"
|
||||
max_tokens = 4096
|
||||
temperature = 0.3
|
||||
|
||||
[[rule]]
|
||||
model_pattern = "qwen3*"
|
||||
provider = "ollama"
|
||||
max_tokens = 2048
|
||||
temperature = 0.3
|
||||
|
||||
[[rule]]
|
||||
model_pattern = "gpt-oss*"
|
||||
provider = "ollama"
|
||||
temperature = 0.1
|
||||
|
||||
[[rule]]
|
||||
model_pattern = "*"
|
||||
provider = "ollama"
|
||||
temperature = 0.5
|
||||
|
||||
# Fallback chain: if primary fails, try these in order
|
||||
fallback = ["ollama", "openrouter"]
|
||||
|
||||
# Cost gating (tokens = cents per 1M)
|
||||
[cost]
|
||||
ollama = 0
|
||||
openrouter = 15
|
||||
claude = 15
|
||||
gemini = 0
|
||||
|
||||
# Daily budget per provider (cents)
|
||||
[daily_budget]
|
||||
ollama = 0
|
||||
openrouter = 1000
|
||||
claude = 500
|
||||
gemini = 0
|
||||
@@ -11,4 +11,3 @@ 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"
|
||||
|
||||
@@ -25,7 +25,7 @@ pub struct EmbedResponse {
|
||||
pub dimensions: usize,
|
||||
}
|
||||
|
||||
#[derive(Clone, Serialize, Deserialize)]
|
||||
#[derive(Serialize, Deserialize)]
|
||||
pub struct GenerateRequest {
|
||||
pub prompt: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
@@ -36,14 +36,6 @@ pub struct GenerateRequest {
|
||||
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)]
|
||||
@@ -142,47 +134,4 @@ impl AiClient {
|
||||
}
|
||||
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}"))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,194 +0,0 @@
|
||||
//! 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;
|
||||
|
||||
/// Rough token count. `chars / 4` ceiling. See module docs for why
|
||||
/// this heuristic is sufficient.
|
||||
pub fn estimate_tokens(text: &str) -> usize {
|
||||
(text.chars().count() + 3) / 4
|
||||
}
|
||||
|
||||
/// 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(estimate_tokens).unwrap_or(0);
|
||||
let estimated = 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::*;
|
||||
|
||||
#[test]
|
||||
fn estimate_tokens_ceiling_divides_by_four() {
|
||||
assert_eq!(estimate_tokens(""), 0);
|
||||
assert_eq!(estimate_tokens("abc"), 1); // 3 → ceil(3/4) = 1
|
||||
assert_eq!(estimate_tokens("abcd"), 1); // 4 → ceil(4/4) = 1
|
||||
assert_eq!(estimate_tokens("abcde"), 2); // 5 → ceil(5/4) = 2
|
||||
assert_eq!(estimate_tokens(&"x".repeat(400)), 100);
|
||||
}
|
||||
|
||||
#[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, 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"));
|
||||
}
|
||||
}
|
||||
@@ -1,461 +0,0 @@
|
||||
//! 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"));
|
||||
}
|
||||
}
|
||||
@@ -1,8 +1,2 @@
|
||||
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;
|
||||
|
||||
@@ -1,39 +0,0 @@
|
||||
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
|
||||
}
|
||||
}
|
||||
@@ -1,2 +0,0 @@
|
||||
pub mod ollama;
|
||||
pub mod openrouter;
|
||||
@@ -1,37 +0,0 @@
|
||||
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
|
||||
}
|
||||
}
|
||||
@@ -1,152 +0,0 @@
|
||||
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",
|
||||
}))
|
||||
}
|
||||
}
|
||||
@@ -1,93 +0,0 @@
|
||||
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('*') {
|
||||
let parts: Vec<&str> = pattern.split('*').collect();
|
||||
if parts.len() == 2 {
|
||||
return name.starts_with(parts[0]) && name.ends_with(parts[1]);
|
||||
} else if parts.len() == 1 {
|
||||
return name.starts_with(parts[0]) || name.ends_with(parts[1]);
|
||||
}
|
||||
}
|
||||
pattern == name
|
||||
}
|
||||
|
||||
impl Default for RoutingRule {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
model_pattern: "*".to_string(),
|
||||
provider: "ollama".to_string(),
|
||||
max_tokens: None,
|
||||
temperature: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -16,17 +16,9 @@ pub fn router(client: AiClient) -> Router {
|
||||
.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)),
|
||||
|
||||
@@ -1,326 +0,0 @@
|
||||
//! 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, estimate_tokens, overflow_message,
|
||||
DEFAULT_MAX_TOKENS, DEFAULT_SAFETY_MARGIN};
|
||||
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 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 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 = 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");
|
||||
}
|
||||
}
|
||||
@@ -15,6 +15,5 @@ bytes = { workspace = true }
|
||||
chrono = { workspace = true }
|
||||
uuid = { workspace = true }
|
||||
object_store = { workspace = true }
|
||||
arrow = { workspace = true }
|
||||
proto = { path = "../proto" }
|
||||
tonic = { workspace = true }
|
||||
|
||||
@@ -54,16 +54,7 @@ impl CatalogService for CatalogGrpc {
|
||||
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)
|
||||
}
|
||||
})?;
|
||||
.map_err(|e| Status::internal(e))?;
|
||||
|
||||
Ok(Response::new(manifest_to_proto(&manifest)))
|
||||
}
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
pub mod registry;
|
||||
pub mod service;
|
||||
pub mod grpc;
|
||||
pub mod tombstones;
|
||||
|
||||
@@ -1,9 +1,7 @@
|
||||
use shared::types::{
|
||||
AiView, ColumnMeta, DatasetId, DatasetManifest, FreshnessContract, Lineage, ModelProfile,
|
||||
ObjectRef, RefreshPolicy, SchemaFingerprint, Sensitivity, Tombstone,
|
||||
DatasetId, DatasetManifest, ObjectRef, SchemaFingerprint,
|
||||
ColumnMeta, Lineage, FreshnessContract, Sensitivity,
|
||||
};
|
||||
|
||||
use crate::tombstones::TombstoneStore;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
@@ -11,44 +9,6 @@ use tokio::sync::RwLock;
|
||||
use storaged::ops;
|
||||
use object_store::ObjectStore;
|
||||
|
||||
/// Make a view name safe for use as an object storage key.
|
||||
/// Allows letters, digits, `_`, `-`, `.`. Anything else becomes `_`.
|
||||
fn sanitize_view_name(name: &str) -> String {
|
||||
name.chars()
|
||||
.map(|c| if c.is_ascii_alphanumeric() || c == '_' || c == '-' || c == '.' { c } else { '_' })
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, serde::Serialize)]
|
||||
pub struct MigrateBucketsReport {
|
||||
pub refs_examined: usize,
|
||||
pub refs_renamed: usize, // legacy "data"/"local" → "primary"
|
||||
pub refs_stamped: usize, // empty → "primary"
|
||||
pub refs_unchanged: usize, // already canonical
|
||||
pub manifests_persisted: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, serde::Serialize)]
|
||||
pub struct DedupeReport {
|
||||
pub groups: usize,
|
||||
pub removed: usize,
|
||||
pub kept: Vec<DedupeKept>,
|
||||
pub errors: Vec<DedupeError>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct DedupeKept {
|
||||
pub name: String,
|
||||
pub kept_id: String,
|
||||
pub removed: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct DedupeError {
|
||||
pub manifest_id: String,
|
||||
pub error: String,
|
||||
}
|
||||
|
||||
/// Partial metadata update — only set fields are applied.
|
||||
#[derive(Debug, Clone, Default, serde::Deserialize)]
|
||||
pub struct MetadataUpdate {
|
||||
@@ -60,23 +20,14 @@ pub struct MetadataUpdate {
|
||||
pub lineage: Option<Lineage>,
|
||||
pub freshness: Option<FreshnessContract>,
|
||||
pub row_count: Option<u64>,
|
||||
// Phase C embedding freshness
|
||||
pub embedding_refresh_policy: Option<RefreshPolicy>,
|
||||
}
|
||||
|
||||
const MANIFEST_PREFIX: &str = "_catalog/manifests";
|
||||
const VIEW_PREFIX: &str = "_catalog/views";
|
||||
const PROFILE_PREFIX: &str = "_catalog/profiles";
|
||||
|
||||
/// In-memory dataset registry backed by manifest persistence in object storage.
|
||||
/// Also tracks AiViews (Phase D) — safe projections over base datasets.
|
||||
/// And tombstones (Phase E) — soft-delete markers applied at query time.
|
||||
#[derive(Clone)]
|
||||
pub struct Registry {
|
||||
datasets: Arc<RwLock<HashMap<DatasetId, DatasetManifest>>>,
|
||||
views: Arc<RwLock<HashMap<String, AiView>>>,
|
||||
profiles: Arc<RwLock<HashMap<String, ModelProfile>>>,
|
||||
tombstones: TombstoneStore,
|
||||
store: Arc<dyn ObjectStore>,
|
||||
}
|
||||
|
||||
@@ -84,47 +35,11 @@ impl Registry {
|
||||
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
|
||||
Self {
|
||||
datasets: Arc::new(RwLock::new(HashMap::new())),
|
||||
views: Arc::new(RwLock::new(HashMap::new())),
|
||||
profiles: Arc::new(RwLock::new(HashMap::new())),
|
||||
tombstones: TombstoneStore::new(store.clone()),
|
||||
store,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn tombstones(&self) -> &TombstoneStore {
|
||||
&self.tombstones
|
||||
}
|
||||
|
||||
/// Add a tombstone for a dataset row. Validates that the dataset
|
||||
/// exists. Idempotent at the endpoint layer — callers dedup if needed.
|
||||
pub async fn add_tombstone(
|
||||
&self,
|
||||
dataset: &str,
|
||||
row_key_column: &str,
|
||||
row_key_value: &str,
|
||||
actor: &str,
|
||||
reason: &str,
|
||||
) -> Result<Tombstone, String> {
|
||||
if self.get_by_name(dataset).await.is_none() {
|
||||
return Err(format!("dataset not found: {dataset}"));
|
||||
}
|
||||
let ts = Tombstone {
|
||||
dataset: dataset.to_string(),
|
||||
row_key_column: row_key_column.to_string(),
|
||||
row_key_value: row_key_value.to_string(),
|
||||
deleted_at: chrono::Utc::now(),
|
||||
actor: actor.to_string(),
|
||||
reason: reason.to_string(),
|
||||
};
|
||||
self.tombstones.append(&ts).await?;
|
||||
Ok(ts)
|
||||
}
|
||||
|
||||
pub async fn list_tombstones(&self, dataset: &str) -> Result<Vec<Tombstone>, String> {
|
||||
self.tombstones.list(dataset).await
|
||||
}
|
||||
|
||||
/// Rebuild in-memory index from persisted manifests + views on startup.
|
||||
/// Rebuild in-memory index from persisted manifests on startup.
|
||||
pub async fn rebuild(&self) -> Result<usize, String> {
|
||||
let keys = ops::list(&self.store, Some(MANIFEST_PREFIX)).await?;
|
||||
let mut datasets = self.datasets.write().await;
|
||||
@@ -137,161 +52,16 @@ impl Registry {
|
||||
}
|
||||
let count = datasets.len();
|
||||
tracing::info!("catalog rebuilt: {count} datasets loaded");
|
||||
|
||||
// Phase D: load AiView definitions alongside manifests.
|
||||
let view_keys = ops::list(&self.store, Some(VIEW_PREFIX)).await.unwrap_or_default();
|
||||
let mut views = self.views.write().await;
|
||||
views.clear();
|
||||
for key in &view_keys {
|
||||
if !key.ends_with(".json") { continue; }
|
||||
let data = match ops::get(&self.store, key).await {
|
||||
Ok(d) => d,
|
||||
Err(e) => {
|
||||
tracing::warn!("view '{key}': read failed: {e}");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
match serde_json::from_slice::<AiView>(&data) {
|
||||
Ok(view) => { views.insert(view.name.clone(), view); }
|
||||
Err(e) => tracing::warn!("view '{key}': parse failed: {e}"),
|
||||
}
|
||||
}
|
||||
if !views.is_empty() {
|
||||
tracing::info!("catalog: {} views loaded", views.len());
|
||||
}
|
||||
|
||||
// Phase 17: load model profiles.
|
||||
let profile_keys = ops::list(&self.store, Some(PROFILE_PREFIX)).await.unwrap_or_default();
|
||||
let mut profiles = self.profiles.write().await;
|
||||
profiles.clear();
|
||||
for key in &profile_keys {
|
||||
if !key.ends_with(".json") { continue; }
|
||||
let data = match ops::get(&self.store, key).await {
|
||||
Ok(d) => d,
|
||||
Err(e) => { tracing::warn!("profile '{key}': read failed: {e}"); continue; }
|
||||
};
|
||||
match serde_json::from_slice::<ModelProfile>(&data) {
|
||||
Ok(p) => { profiles.insert(p.id.clone(), p); }
|
||||
Err(e) => tracing::warn!("profile '{key}': parse failed: {e}"),
|
||||
}
|
||||
}
|
||||
if !profiles.is_empty() {
|
||||
tracing::info!("catalog: {} model profiles loaded", profiles.len());
|
||||
}
|
||||
|
||||
Ok(count)
|
||||
}
|
||||
|
||||
// --- Phase 17: Model profiles ---
|
||||
|
||||
/// Create or replace a model profile. Validates id slug + ensures
|
||||
/// every bound_dataset exists (as either a raw dataset or an AiView).
|
||||
pub async fn put_profile(&self, profile: ModelProfile) -> Result<ModelProfile, String> {
|
||||
if profile.id.is_empty() {
|
||||
return Err("profile id is empty".into());
|
||||
}
|
||||
if !profile.id.chars().all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_') {
|
||||
return Err(format!(
|
||||
"profile id '{}' must be alphanumeric + '-' or '_' only", profile.id,
|
||||
));
|
||||
}
|
||||
for binding in &profile.bound_datasets {
|
||||
let exists = self.get_by_name(binding).await.is_some()
|
||||
|| self.get_view(binding).await.is_some();
|
||||
if !exists {
|
||||
return Err(format!(
|
||||
"bound dataset '{}' not found as dataset or view", binding,
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
let key = format!("{PROFILE_PREFIX}/{}.json", profile.id);
|
||||
let json = serde_json::to_vec_pretty(&profile).map_err(|e| e.to_string())?;
|
||||
ops::put(&self.store, &key, json.into()).await?;
|
||||
|
||||
let mut profiles = self.profiles.write().await;
|
||||
profiles.insert(profile.id.clone(), profile.clone());
|
||||
tracing::info!(
|
||||
"profile registered: {} -> ollama={} bindings={:?}",
|
||||
profile.id, profile.ollama_name, profile.bound_datasets,
|
||||
);
|
||||
Ok(profile)
|
||||
}
|
||||
|
||||
pub async fn get_profile(&self, id: &str) -> Option<ModelProfile> {
|
||||
self.profiles.read().await.get(id).cloned()
|
||||
}
|
||||
|
||||
pub async fn list_profiles(&self) -> Vec<ModelProfile> {
|
||||
self.profiles.read().await.values().cloned().collect()
|
||||
}
|
||||
|
||||
pub async fn delete_profile(&self, id: &str) -> Result<(), String> {
|
||||
let key = format!("{PROFILE_PREFIX}/{id}.json");
|
||||
ops::delete(&self.store, &key).await?;
|
||||
self.profiles.write().await.remove(id);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Register a dataset. Idempotent on `name`:
|
||||
/// - No existing dataset by this name → create new manifest.
|
||||
/// - Exists with same schema fingerprint → update objects + `updated_at`
|
||||
/// in place (re-ingest of identical-shape data).
|
||||
/// - Exists with different schema fingerprint → reject as `Err`
|
||||
/// (callers map to HTTP 409 / gRPC FAILED_PRECONDITION). Schema
|
||||
/// evolution is not yet exposed as an HTTP endpoint; callers must
|
||||
/// register under a new name or delete the existing manifest first.
|
||||
///
|
||||
/// Concurrency: the write lock is held across the storage write. That
|
||||
/// serializes concurrent registers but is correctness-mandatory —
|
||||
/// dropping the lock across the check→insert sequence creates a TOCTOU
|
||||
/// window where two callers with the same name both see "no existing"
|
||||
/// and both insert, reintroducing the duplicate-manifest bug this
|
||||
/// function exists to prevent. Register is a metadata-only hop on the
|
||||
/// ingest path, so the serialization cost is acceptable. Precedent:
|
||||
/// `update_metadata` also holds the write lock across its I/O.
|
||||
///
|
||||
/// Legacy state: if multiple manifests with the same name already exist
|
||||
/// (from before this was idempotent), the most recently updated one is
|
||||
/// treated as canonical. Run `dedupe_by_name` to clean up the rest.
|
||||
/// Register a new dataset. Persists manifest to storage before updating memory.
|
||||
pub async fn register(
|
||||
&self,
|
||||
name: String,
|
||||
schema_fingerprint: SchemaFingerprint,
|
||||
objects: Vec<ObjectRef>,
|
||||
) -> Result<DatasetManifest, String> {
|
||||
let mut datasets = self.datasets.write().await;
|
||||
|
||||
let existing = datasets
|
||||
.values()
|
||||
.filter(|d| d.name == name)
|
||||
.max_by_key(|d| d.updated_at)
|
||||
.cloned();
|
||||
|
||||
if let Some(mut manifest) = existing {
|
||||
if manifest.schema_fingerprint != schema_fingerprint {
|
||||
return Err(format!(
|
||||
"dataset '{}' already exists with a different schema \
|
||||
(existing fingerprint: {}, new: {}). Schema-evolution \
|
||||
migration is not yet exposed as an HTTP endpoint; \
|
||||
register under a new name or delete the existing \
|
||||
manifest first.",
|
||||
name, manifest.schema_fingerprint.0, schema_fingerprint.0,
|
||||
));
|
||||
}
|
||||
manifest.objects = objects;
|
||||
manifest.updated_at = chrono::Utc::now();
|
||||
|
||||
let manifest_key = format!("{MANIFEST_PREFIX}/{}.json", manifest.id);
|
||||
let json = serde_json::to_vec_pretty(&manifest).map_err(|e| e.to_string())?;
|
||||
ops::put(&self.store, &manifest_key, json.into()).await?;
|
||||
|
||||
datasets.insert(manifest.id.clone(), manifest.clone());
|
||||
|
||||
tracing::info!("re-registered (idempotent): {} ({})", manifest.name, manifest.id);
|
||||
return Ok(manifest);
|
||||
}
|
||||
|
||||
let now = chrono::Utc::now();
|
||||
let manifest = DatasetManifest {
|
||||
id: DatasetId::new(),
|
||||
@@ -308,15 +78,14 @@ impl Registry {
|
||||
freshness: None,
|
||||
tags: vec![],
|
||||
row_count: None,
|
||||
last_embedded_at: None,
|
||||
embedding_stale_since: None,
|
||||
embedding_refresh_policy: None,
|
||||
};
|
||||
|
||||
// Write-ahead: persist before in-memory update
|
||||
let manifest_key = format!("{MANIFEST_PREFIX}/{}.json", manifest.id);
|
||||
let json = serde_json::to_vec_pretty(&manifest).map_err(|e| e.to_string())?;
|
||||
ops::put(&self.store, &manifest_key, json.into()).await?;
|
||||
|
||||
let mut datasets = self.datasets.write().await;
|
||||
datasets.insert(manifest.id.clone(), manifest.clone());
|
||||
|
||||
tracing::info!("registered dataset: {} ({})", manifest.name, manifest.id);
|
||||
@@ -342,7 +111,6 @@ impl Registry {
|
||||
if let Some(lineage) = updates.lineage { manifest.lineage = Some(lineage); }
|
||||
if let Some(freshness) = updates.freshness { manifest.freshness = Some(freshness); }
|
||||
if let Some(count) = updates.row_count { manifest.row_count = Some(count); }
|
||||
if let Some(policy) = updates.embedding_refresh_policy { manifest.embedding_refresh_policy = Some(policy); }
|
||||
manifest.updated_at = chrono::Utc::now();
|
||||
|
||||
// Persist
|
||||
@@ -372,370 +140,6 @@ impl Registry {
|
||||
datasets.values().cloned().collect()
|
||||
}
|
||||
|
||||
/// Re-read the parquet footer(s) for a dataset and repopulate `row_count`
|
||||
/// and `columns` from reality. Use this to repair manifests whose
|
||||
/// metadata was lost (e.g. migrated from a pre-Phase 10 catalog).
|
||||
///
|
||||
/// Does NOT touch owner/description/sensitivity/lineage/tags — only
|
||||
/// the structural facts that parquet can tell us authoritatively.
|
||||
/// The existing `schema_fingerprint` is updated if the recomputed one
|
||||
/// differs; a warning is logged so drift is visible.
|
||||
pub async fn resync_from_parquet(&self, name: &str) -> Result<DatasetManifest, String> {
|
||||
use shared::arrow_helpers::{fingerprint_schema, parquet_to_record_batches};
|
||||
|
||||
// Snapshot the target manifest so we don't hold the write lock during IO.
|
||||
let (id, objects, old_fp) = {
|
||||
let datasets = self.datasets.read().await;
|
||||
let m = datasets
|
||||
.values()
|
||||
.find(|d| d.name == name)
|
||||
.ok_or_else(|| format!("dataset not found: {name}"))?;
|
||||
(m.id.clone(), m.objects.clone(), m.schema_fingerprint.clone())
|
||||
};
|
||||
|
||||
if objects.is_empty() {
|
||||
return Err(format!("dataset '{name}' has no object references to resync from"));
|
||||
}
|
||||
|
||||
let mut total_rows: u64 = 0;
|
||||
let mut first_schema: Option<arrow::datatypes::SchemaRef> = None;
|
||||
for obj in &objects {
|
||||
let data = ops::get(&self.store, &obj.key).await
|
||||
.map_err(|e| format!("read {}: {e}", obj.key))?;
|
||||
let (schema, batches) = parquet_to_record_batches(&data)
|
||||
.map_err(|e| format!("parse {}: {e}", obj.key))?;
|
||||
let rows: u64 = batches.iter().map(|b| b.num_rows() as u64).sum();
|
||||
total_rows += rows;
|
||||
if first_schema.is_none() {
|
||||
first_schema = Some(schema);
|
||||
}
|
||||
}
|
||||
|
||||
let schema = first_schema.ok_or("no schema recovered")?;
|
||||
let new_fp = fingerprint_schema(&schema);
|
||||
if new_fp != old_fp {
|
||||
tracing::warn!(
|
||||
"dataset '{}' schema fingerprint drift: {} -> {} (updating to match parquet reality)",
|
||||
name, old_fp.0, new_fp.0,
|
||||
);
|
||||
}
|
||||
|
||||
let columns: Vec<ColumnMeta> = schema
|
||||
.fields()
|
||||
.iter()
|
||||
.map(|f| ColumnMeta {
|
||||
name: f.name().clone(),
|
||||
data_type: f.data_type().to_string(),
|
||||
sensitivity: None,
|
||||
description: String::new(),
|
||||
is_pii: false,
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Apply updates.
|
||||
let mut datasets = self.datasets.write().await;
|
||||
let manifest = datasets
|
||||
.get_mut(&id)
|
||||
.ok_or_else(|| format!("dataset disappeared during resync: {name}"))?;
|
||||
manifest.row_count = Some(total_rows);
|
||||
manifest.columns = columns;
|
||||
manifest.schema_fingerprint = new_fp;
|
||||
manifest.updated_at = chrono::Utc::now();
|
||||
|
||||
// Persist.
|
||||
let manifest_key = format!("{MANIFEST_PREFIX}/{}.json", manifest.id);
|
||||
let json = serde_json::to_vec_pretty(manifest).map_err(|e| e.to_string())?;
|
||||
ops::put(&self.store, &manifest_key, json.into()).await?;
|
||||
|
||||
tracing::info!("resynced '{name}': row_count={total_rows}, {} columns", manifest.columns.len());
|
||||
Ok(manifest.clone())
|
||||
}
|
||||
|
||||
/// Collapse duplicate manifests that share a `name`. Winner per group:
|
||||
/// 1. Prefer a manifest with a non-null `row_count` (already resynced).
|
||||
/// 2. Break ties by newest `updated_at`.
|
||||
/// Losers are removed from the in-memory registry AND their manifest
|
||||
/// JSON is deleted from object storage. Datasets with a single manifest
|
||||
/// are untouched. Parquet data files are never touched — only catalog
|
||||
/// metadata. See `register` for the prevention side of this fix.
|
||||
///
|
||||
/// Operator-only. Do not run while ingest is active: a concurrent
|
||||
/// `register` between the snapshot and the delete sweep can create a
|
||||
/// new manifest under a name we just deduped, producing a transient
|
||||
/// count of 2 until the next dedupe run. Safe outside ingest windows.
|
||||
pub async fn dedupe_by_name(&self) -> DedupeReport {
|
||||
let groups: Vec<(String, Vec<DatasetManifest>)> = {
|
||||
let datasets = self.datasets.read().await;
|
||||
let mut by_name: HashMap<String, Vec<DatasetManifest>> = HashMap::new();
|
||||
for m in datasets.values() {
|
||||
by_name.entry(m.name.clone()).or_default().push(m.clone());
|
||||
}
|
||||
by_name.into_iter().filter(|(_, v)| v.len() > 1).collect()
|
||||
};
|
||||
|
||||
let mut report = DedupeReport::default();
|
||||
let mut to_delete: Vec<DatasetId> = Vec::new();
|
||||
|
||||
for (name, mut manifests) in groups {
|
||||
report.groups += 1;
|
||||
manifests.sort_by(|a, b| {
|
||||
b.row_count.is_some().cmp(&a.row_count.is_some())
|
||||
.then(b.updated_at.cmp(&a.updated_at))
|
||||
});
|
||||
let winner = &manifests[0];
|
||||
report.kept.push(DedupeKept {
|
||||
name: name.clone(),
|
||||
kept_id: winner.id.to_string(),
|
||||
removed: manifests.len() - 1,
|
||||
});
|
||||
for loser in &manifests[1..] {
|
||||
to_delete.push(loser.id.clone());
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
let mut datasets = self.datasets.write().await;
|
||||
for id in &to_delete {
|
||||
datasets.remove(id);
|
||||
}
|
||||
}
|
||||
for id in &to_delete {
|
||||
let key = format!("{MANIFEST_PREFIX}/{}.json", id);
|
||||
match ops::delete(&self.store, &key).await {
|
||||
Ok(_) => report.removed += 1,
|
||||
Err(e) => report.errors.push(DedupeError {
|
||||
manifest_id: id.to_string(),
|
||||
error: e,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
"dedupe: {} groups collapsed, {} manifests removed, {} errors",
|
||||
report.groups, report.removed, report.errors.len(),
|
||||
);
|
||||
report
|
||||
}
|
||||
|
||||
/// Resync every dataset that currently has a null row_count.
|
||||
/// Returns (successes, failures) where each entry is (name, detail).
|
||||
pub async fn resync_missing(&self) -> (Vec<(String, u64)>, Vec<(String, String)>) {
|
||||
let names: Vec<String> = {
|
||||
let datasets = self.datasets.read().await;
|
||||
datasets
|
||||
.values()
|
||||
.filter(|d| d.row_count.is_none() || d.columns.is_empty())
|
||||
.map(|d| d.name.clone())
|
||||
.collect()
|
||||
};
|
||||
|
||||
let mut ok = Vec::new();
|
||||
let mut err = Vec::new();
|
||||
for name in names {
|
||||
match self.resync_from_parquet(&name).await {
|
||||
Ok(m) => ok.push((name, m.row_count.unwrap_or(0))),
|
||||
Err(e) => err.push((name, e)),
|
||||
}
|
||||
}
|
||||
(ok, err)
|
||||
}
|
||||
|
||||
/// Mark a dataset's embeddings as stale (row-level data has been written
|
||||
/// since the last embedding refresh). Idempotent — setting stale when
|
||||
/// already stale is a no-op. Only marks stale if the dataset has been
|
||||
/// embedded before — a never-embedded dataset doesn't need a stale flag
|
||||
/// (it just needs an initial index build). Called from the ingest path.
|
||||
pub async fn mark_embeddings_stale(&self, name: &str) -> Result<(), String> {
|
||||
let mut datasets = self.datasets.write().await;
|
||||
let manifest = datasets
|
||||
.values_mut()
|
||||
.find(|d| d.name == name)
|
||||
.ok_or_else(|| format!("dataset not found: {name}"))?;
|
||||
|
||||
if manifest.last_embedded_at.is_none() {
|
||||
return Ok(()); // never embedded -> no stale semantics yet
|
||||
}
|
||||
if manifest.embedding_stale_since.is_none() {
|
||||
manifest.embedding_stale_since = Some(chrono::Utc::now());
|
||||
manifest.updated_at = chrono::Utc::now();
|
||||
|
||||
let key = format!("{MANIFEST_PREFIX}/{}.json", manifest.id);
|
||||
let json = serde_json::to_vec_pretty(manifest).map_err(|e| e.to_string())?;
|
||||
ops::put(&self.store, &key, json.into()).await?;
|
||||
tracing::info!("marked embeddings stale for dataset '{name}'");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Clear the stale marker and set `last_embedded_at = now`.
|
||||
/// Called by the embedding refresh pipeline once it finishes.
|
||||
pub async fn clear_embeddings_stale(&self, name: &str) -> Result<(), String> {
|
||||
let mut datasets = self.datasets.write().await;
|
||||
let manifest = datasets
|
||||
.values_mut()
|
||||
.find(|d| d.name == name)
|
||||
.ok_or_else(|| format!("dataset not found: {name}"))?;
|
||||
|
||||
let now = chrono::Utc::now();
|
||||
manifest.embedding_stale_since = None;
|
||||
manifest.last_embedded_at = Some(now);
|
||||
manifest.updated_at = now;
|
||||
|
||||
let key = format!("{MANIFEST_PREFIX}/{}.json", manifest.id);
|
||||
let json = serde_json::to_vec_pretty(manifest).map_err(|e| e.to_string())?;
|
||||
ops::put(&self.store, &key, json.into()).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Federation layer 2: stamp `bucket = "primary"` on every ObjectRef
|
||||
/// whose `bucket` field is empty or matches a legacy value (`"data"`,
|
||||
/// `"local"`). One-shot migration; re-running is a safe no-op once
|
||||
/// every ref is canonical.
|
||||
pub async fn migrate_buckets_to_primary(&self) -> Result<MigrateBucketsReport, String> {
|
||||
let mut datasets = self.datasets.write().await;
|
||||
let mut report = MigrateBucketsReport::default();
|
||||
let mut to_persist: Vec<(DatasetId, DatasetManifest)> = Vec::new();
|
||||
|
||||
for manifest in datasets.values_mut() {
|
||||
let mut changed = false;
|
||||
for obj in manifest.objects.iter_mut() {
|
||||
report.refs_examined += 1;
|
||||
if obj.bucket.is_empty() {
|
||||
obj.bucket = "primary".to_string();
|
||||
report.refs_stamped += 1;
|
||||
changed = true;
|
||||
} else if obj.bucket == "data" || obj.bucket == "local" {
|
||||
obj.bucket = "primary".to_string();
|
||||
report.refs_renamed += 1;
|
||||
changed = true;
|
||||
} else {
|
||||
report.refs_unchanged += 1;
|
||||
}
|
||||
}
|
||||
if changed {
|
||||
manifest.updated_at = chrono::Utc::now();
|
||||
to_persist.push((manifest.id.clone(), manifest.clone()));
|
||||
}
|
||||
}
|
||||
|
||||
// Persist updated manifests after we've finished mutating the map.
|
||||
for (id, manifest) in to_persist {
|
||||
let key = format!("{MANIFEST_PREFIX}/{}.json", id);
|
||||
let json = serde_json::to_vec_pretty(&manifest).map_err(|e| e.to_string())?;
|
||||
ops::put(&self.store, &key, json.into()).await?;
|
||||
report.manifests_persisted += 1;
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
"bucket migration: examined {} refs, renamed {}, stamped {}, unchanged {}, persisted {} manifests",
|
||||
report.refs_examined,
|
||||
report.refs_renamed,
|
||||
report.refs_stamped,
|
||||
report.refs_unchanged,
|
||||
report.manifests_persisted,
|
||||
);
|
||||
Ok(report)
|
||||
}
|
||||
|
||||
// --- Phase D: AI-safe views ---
|
||||
|
||||
/// Create or replace a named view. Validates that the base dataset
|
||||
/// exists and that the column whitelist is non-empty. Persists to
|
||||
/// `_catalog/views/{name}.json`.
|
||||
pub async fn put_view(&self, mut view: AiView) -> Result<AiView, String> {
|
||||
if view.name.is_empty() {
|
||||
return Err("view name is empty".into());
|
||||
}
|
||||
if view.columns.is_empty() {
|
||||
return Err("view must whitelist at least one column".into());
|
||||
}
|
||||
// Base dataset must exist (read-side will fail anyway, fail-fast here
|
||||
// so operators don't end up with dangling views).
|
||||
if self.get_by_name(&view.base_dataset).await.is_none() {
|
||||
return Err(format!(
|
||||
"base dataset '{}' not found in catalog",
|
||||
view.base_dataset
|
||||
));
|
||||
}
|
||||
if view.created_at.timestamp() == 0 {
|
||||
view.created_at = chrono::Utc::now();
|
||||
}
|
||||
|
||||
let key = format!("{VIEW_PREFIX}/{}.json", sanitize_view_name(&view.name));
|
||||
let json = serde_json::to_vec_pretty(&view).map_err(|e| e.to_string())?;
|
||||
ops::put(&self.store, &key, json.into()).await?;
|
||||
|
||||
let mut views = self.views.write().await;
|
||||
views.insert(view.name.clone(), view.clone());
|
||||
tracing::info!("view registered: {} over '{}'", view.name, view.base_dataset);
|
||||
Ok(view)
|
||||
}
|
||||
|
||||
pub async fn get_view(&self, name: &str) -> Option<AiView> {
|
||||
self.views.read().await.get(name).cloned()
|
||||
}
|
||||
|
||||
pub async fn list_views(&self) -> Vec<AiView> {
|
||||
self.views.read().await.values().cloned().collect()
|
||||
}
|
||||
|
||||
pub async fn delete_view(&self, name: &str) -> Result<(), String> {
|
||||
let key = format!("{VIEW_PREFIX}/{}.json", sanitize_view_name(name));
|
||||
ops::delete(&self.store, &key).await?;
|
||||
self.views.write().await.remove(name);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Remove a dataset from the catalog by name. Deletes the manifest
|
||||
/// from both the in-memory registry and object storage.
|
||||
///
|
||||
/// Scope is metadata-only: the underlying parquet files, vector
|
||||
/// indexes, tombstones, trial journals, and AiViews that reference
|
||||
/// this dataset are NOT touched. Caller is responsible for any
|
||||
/// cascade cleanup. This mirrors how a DROP TABLE in a typical
|
||||
/// warehouse separates "forget about this dataset" from "actually
|
||||
/// reclaim the bytes".
|
||||
///
|
||||
/// Returns Ok(()) when the dataset existed and was removed; an Err
|
||||
/// with "dataset not found" when no manifest by that name exists.
|
||||
/// Legacy state with duplicate names: all matching manifests are
|
||||
/// removed (effectively a post-hoc dedupe on that name).
|
||||
pub async fn delete_dataset(&self, name: &str) -> Result<usize, String> {
|
||||
let ids_to_remove: Vec<DatasetId> = {
|
||||
let datasets = self.datasets.read().await;
|
||||
datasets.values().filter(|d| d.name == name).map(|d| d.id.clone()).collect()
|
||||
};
|
||||
if ids_to_remove.is_empty() {
|
||||
return Err(format!("dataset not found: {name}"));
|
||||
}
|
||||
|
||||
for id in &ids_to_remove {
|
||||
let key = format!("{MANIFEST_PREFIX}/{}.json", id);
|
||||
if let Err(e) = ops::delete(&self.store, &key).await {
|
||||
// Storage delete failed — log and keep going. The in-memory
|
||||
// remove below is still correct; on restart, the missing
|
||||
// storage object just means nothing to rehydrate.
|
||||
tracing::warn!("delete_dataset '{name}': storage delete of {key} failed: {e}");
|
||||
}
|
||||
}
|
||||
let mut datasets = self.datasets.write().await;
|
||||
for id in &ids_to_remove {
|
||||
datasets.remove(id);
|
||||
}
|
||||
tracing::info!("deleted dataset '{}' ({} manifest(s))", name, ids_to_remove.len());
|
||||
Ok(ids_to_remove.len())
|
||||
}
|
||||
|
||||
/// List datasets whose `embedding_stale_since` is set — they need a refresh.
|
||||
pub async fn stale_datasets(&self) -> Vec<DatasetManifest> {
|
||||
let datasets = self.datasets.read().await;
|
||||
datasets
|
||||
.values()
|
||||
.filter(|d| d.embedding_stale_since.is_some())
|
||||
.cloned()
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Add objects to an existing dataset.
|
||||
pub async fn add_objects(
|
||||
&self,
|
||||
@@ -756,275 +160,3 @@ impl Registry {
|
||||
Ok(manifest.clone())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use object_store::memory::InMemory;
|
||||
|
||||
fn fixture() -> Registry {
|
||||
Registry::new(Arc::new(InMemory::new()))
|
||||
}
|
||||
|
||||
fn fp(s: &str) -> SchemaFingerprint {
|
||||
SchemaFingerprint(s.to_string())
|
||||
}
|
||||
|
||||
fn obj(key: &str) -> ObjectRef {
|
||||
ObjectRef {
|
||||
bucket: "primary".to_string(),
|
||||
key: key.to_string(),
|
||||
size_bytes: 100,
|
||||
created_at: chrono::Utc::now(),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fresh_register_creates_new_manifest() {
|
||||
let reg = fixture();
|
||||
let m = reg.register("a".into(), fp("f1"), vec![obj("a.parquet")]).await.unwrap();
|
||||
assert_eq!(m.name, "a");
|
||||
assert_eq!(m.schema_fingerprint, fp("f1"));
|
||||
assert_eq!(reg.list().await.len(), 1);
|
||||
assert_eq!(reg.get_by_name("a").await.unwrap().id, m.id);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn re_register_same_fingerprint_is_idempotent() {
|
||||
let reg = fixture();
|
||||
let first = reg.register("a".into(), fp("f1"), vec![obj("old.parquet")]).await.unwrap();
|
||||
let second = reg.register("a".into(), fp("f1"), vec![obj("new.parquet")]).await.unwrap();
|
||||
|
||||
assert_eq!(first.id, second.id, "same name+fp should reuse ID");
|
||||
assert_eq!(reg.list().await.len(), 1, "no duplicate manifest created");
|
||||
let fetched = reg.get_by_name("a").await.unwrap();
|
||||
assert_eq!(fetched.objects[0].key, "new.parquet", "objects replaced");
|
||||
assert!(fetched.updated_at >= first.updated_at, "updated_at bumped");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn re_register_different_fingerprint_rejects() {
|
||||
let reg = fixture();
|
||||
reg.register("a".into(), fp("f1"), vec![obj("a.parquet")]).await.unwrap();
|
||||
let err = reg.register("a".into(), fp("f2"), vec![obj("a.parquet")]).await.unwrap_err();
|
||||
|
||||
assert!(err.contains("different schema"), "error mentions schema drift: {err}");
|
||||
assert_eq!(reg.list().await.len(), 1, "failed register did not mutate state");
|
||||
assert_eq!(reg.get_by_name("a").await.unwrap().schema_fingerprint, fp("f1"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn delete_dataset_removes_manifest_and_in_memory_entry() {
|
||||
let reg = fixture();
|
||||
reg.register("to_delete".into(), fp("f1"), vec![obj("d.parquet")]).await.unwrap();
|
||||
reg.register("keepme".into(), fp("f2"), vec![obj("k.parquet")]).await.unwrap();
|
||||
assert_eq!(reg.list().await.len(), 2);
|
||||
|
||||
let removed = reg.delete_dataset("to_delete").await.unwrap();
|
||||
assert_eq!(removed, 1, "one manifest removed");
|
||||
assert!(reg.get_by_name("to_delete").await.is_none());
|
||||
assert!(reg.get_by_name("keepme").await.is_some(), "unrelated dataset untouched");
|
||||
assert_eq!(reg.list().await.len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn delete_dataset_404s_on_unknown_name() {
|
||||
let reg = fixture();
|
||||
let err = reg.delete_dataset("nope").await.unwrap_err();
|
||||
assert!(err.starts_with("dataset not found"), "error wording: {err}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn delete_dataset_removes_all_dupes_with_matching_name() {
|
||||
// Legacy zombies: direct insert so we can seed duplicates,
|
||||
// mirroring pre-idempotent-register state. delete_dataset should
|
||||
// sweep every manifest sharing that name in a single call.
|
||||
let reg = fixture();
|
||||
let now = chrono::Utc::now();
|
||||
for _ in 0..3 {
|
||||
let m = DatasetManifest {
|
||||
id: DatasetId::new(),
|
||||
name: "zombie".into(),
|
||||
schema_fingerprint: fp("fp"),
|
||||
objects: vec![obj("z.parquet")],
|
||||
created_at: now, updated_at: now,
|
||||
row_count: None, description: String::new(), owner: String::new(),
|
||||
sensitivity: None, columns: vec![], lineage: None, freshness: None,
|
||||
tags: vec![], last_embedded_at: None, embedding_stale_since: None,
|
||||
embedding_refresh_policy: None,
|
||||
};
|
||||
reg.datasets.write().await.insert(m.id.clone(), m);
|
||||
}
|
||||
assert_eq!(reg.list().await.len(), 3);
|
||||
|
||||
let removed = reg.delete_dataset("zombie").await.unwrap();
|
||||
assert_eq!(removed, 3, "all three manifests removed in one sweep");
|
||||
assert_eq!(reg.list().await.len(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dedupe_no_duplicates_is_noop() {
|
||||
let reg = fixture();
|
||||
reg.register("a".into(), fp("f1"), vec![obj("a.parquet")]).await.unwrap();
|
||||
reg.register("b".into(), fp("f2"), vec![obj("b.parquet")]).await.unwrap();
|
||||
|
||||
let report = reg.dedupe_by_name().await;
|
||||
assert_eq!(report.groups, 0);
|
||||
assert_eq!(report.removed, 0);
|
||||
assert_eq!(reg.list().await.len(), 2);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 8)]
|
||||
async fn concurrent_register_stress_many_workers_one_manifest() {
|
||||
// Stress variant: fire 32 concurrent registers for the same new
|
||||
// name across 8 worker threads. Single-manifest invariant must
|
||||
// hold regardless of how the scheduler interleaves them. This
|
||||
// catches race regressions the 2-call variant below might miss
|
||||
// under a forgiving scheduler.
|
||||
let reg = fixture();
|
||||
let mut handles = Vec::new();
|
||||
for i in 0..32 {
|
||||
let r = reg.clone();
|
||||
handles.push(tokio::spawn(async move {
|
||||
r.register(
|
||||
"stress".into(),
|
||||
fp("stress-fp"),
|
||||
vec![obj(&format!("{i}.parquet"))],
|
||||
).await
|
||||
}));
|
||||
}
|
||||
for h in handles {
|
||||
h.await.unwrap().expect("each register succeeds (idempotent)");
|
||||
}
|
||||
let all = reg.list().await;
|
||||
let stress_manifests: Vec<_> = all.iter().filter(|d| d.name == "stress").collect();
|
||||
assert_eq!(stress_manifests.len(), 1,
|
||||
"32 concurrent registers produced {} manifests for 'stress' — expected 1",
|
||||
stress_manifests.len());
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
async fn concurrent_register_same_new_name_collapses_to_one_manifest() {
|
||||
// Pins the TOCTOU fix: without the write-lock held across the
|
||||
// check→insert sequence, two parallel registers of the same
|
||||
// previously-unknown name would both see "no existing" and both
|
||||
// insert. Multi-thread flavor is required to actually exercise
|
||||
// the race — under single-threaded tokio the scheduler doesn't
|
||||
// preempt between awaits on the same task's critical section.
|
||||
// Under the current impl, whichever acquires the write lock
|
||||
// first creates the manifest; the second observes it and takes
|
||||
// the idempotent-update branch.
|
||||
let reg = fixture();
|
||||
let a = reg.clone();
|
||||
let b = reg.clone();
|
||||
|
||||
let (ra, rb) = tokio::join!(
|
||||
async move {
|
||||
a.register("race".into(), fp("shared-fp"), vec![obj("a.parquet")]).await
|
||||
},
|
||||
async move {
|
||||
b.register("race".into(), fp("shared-fp"), vec![obj("b.parquet")]).await
|
||||
},
|
||||
);
|
||||
|
||||
let m_a = ra.expect("first register succeeds");
|
||||
let m_b = rb.expect("second register succeeds (same fp = idempotent)");
|
||||
assert_eq!(m_a.id, m_b.id, "both calls resolve to a single DatasetId");
|
||||
assert_eq!(reg.list().await.len(), 1, "no duplicate manifest created");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
async fn concurrent_register_same_name_different_fp_one_wins_one_rejects() {
|
||||
// Reverse case: concurrent registers with conflicting fingerprints.
|
||||
// Deterministic outcome: one call's fingerprint gets persisted as
|
||||
// canonical, the other is rejected on seeing the now-established
|
||||
// fingerprint. Which one wins is scheduler-dependent — we only
|
||||
// assert exactly one success and exactly one rejection.
|
||||
let reg = fixture();
|
||||
let a = reg.clone();
|
||||
let b = reg.clone();
|
||||
|
||||
let (ra, rb) = tokio::join!(
|
||||
async move { a.register("race".into(), fp("fp1"), vec![obj("a.parquet")]).await },
|
||||
async move { b.register("race".into(), fp("fp2"), vec![obj("b.parquet")]).await },
|
||||
);
|
||||
|
||||
let successes = [&ra, &rb].iter().filter(|r| r.is_ok()).count();
|
||||
let rejections = [&ra, &rb].iter().filter(|r| r.is_err()).count();
|
||||
assert_eq!(successes, 1, "exactly one register wins");
|
||||
assert_eq!(rejections, 1, "the other is rejected on fingerprint");
|
||||
let rejection_msg = [&ra, &rb].iter()
|
||||
.find_map(|r| r.as_ref().err())
|
||||
.unwrap();
|
||||
assert!(rejection_msg.contains("different schema"));
|
||||
assert_eq!(reg.list().await.len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dedupe_collapses_dupes_preferring_non_null_row_count() {
|
||||
let reg = fixture();
|
||||
|
||||
// Simulate legacy pre-idempotent-register state by inserting
|
||||
// manifests with the same `name` directly. Mirrors what happened
|
||||
// with the 308× successful_playbooks duplication.
|
||||
let now = chrono::Utc::now();
|
||||
let old_with_rows = DatasetManifest {
|
||||
id: DatasetId::new(),
|
||||
name: "dupes".into(),
|
||||
schema_fingerprint: fp("f1"),
|
||||
objects: vec![obj("dupes.parquet")],
|
||||
created_at: now - chrono::Duration::hours(2),
|
||||
updated_at: now - chrono::Duration::hours(2),
|
||||
row_count: Some(42),
|
||||
description: String::new(),
|
||||
owner: String::new(),
|
||||
sensitivity: None,
|
||||
columns: vec![],
|
||||
lineage: None,
|
||||
freshness: None,
|
||||
tags: vec![],
|
||||
last_embedded_at: None,
|
||||
embedding_stale_since: None,
|
||||
embedding_refresh_policy: None,
|
||||
};
|
||||
let newer_no_rows = DatasetManifest {
|
||||
id: DatasetId::new(),
|
||||
updated_at: now, // newer, but null row_count
|
||||
row_count: None,
|
||||
..old_with_rows.clone()
|
||||
};
|
||||
let oldest = DatasetManifest {
|
||||
id: DatasetId::new(),
|
||||
updated_at: now - chrono::Duration::hours(5),
|
||||
row_count: None,
|
||||
..old_with_rows.clone()
|
||||
};
|
||||
|
||||
let kept_id = old_with_rows.id.clone();
|
||||
let loser_a = newer_no_rows.id.clone();
|
||||
let loser_b = oldest.id.clone();
|
||||
|
||||
{
|
||||
let mut d = reg.datasets.write().await;
|
||||
d.insert(old_with_rows.id.clone(), old_with_rows);
|
||||
d.insert(newer_no_rows.id.clone(), newer_no_rows);
|
||||
d.insert(oldest.id.clone(), oldest);
|
||||
}
|
||||
assert_eq!(reg.list().await.len(), 3);
|
||||
|
||||
let report = reg.dedupe_by_name().await;
|
||||
assert_eq!(report.groups, 1);
|
||||
assert_eq!(report.removed, 2);
|
||||
assert_eq!(report.errors.len(), 0);
|
||||
assert_eq!(report.kept.len(), 1);
|
||||
assert_eq!(report.kept[0].name, "dupes");
|
||||
assert_eq!(report.kept[0].kept_id, kept_id.to_string(),
|
||||
"should keep the manifest with non-null row_count");
|
||||
|
||||
let remaining = reg.list().await;
|
||||
assert_eq!(remaining.len(), 1);
|
||||
assert_eq!(remaining[0].id, kept_id);
|
||||
assert!(reg.get(&loser_a).await.is_none());
|
||||
assert!(reg.get(&loser_b).await.is_none());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -17,24 +17,8 @@ pub fn router(registry: Registry) -> Router {
|
||||
.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}", get(get_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)
|
||||
}
|
||||
|
||||
@@ -126,14 +110,7 @@ async fn create_dataset(
|
||||
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))
|
||||
}
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -165,25 +142,6 @@ async fn get_dataset_by_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>,
|
||||
@@ -194,282 +152,3 @@ async fn update_metadata(
|
||||
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)),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,151 +0,0 @@
|
||||
//! 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)
|
||||
}
|
||||
}
|
||||
@@ -12,7 +12,6 @@ aibridge = { path = "../aibridge" }
|
||||
ingestd = { path = "../ingestd" }
|
||||
vectord = { path = "../vectord" }
|
||||
journald = { path = "../journald" }
|
||||
truth = { path = "../truth" }
|
||||
tokio = { workspace = true }
|
||||
axum = { workspace = true }
|
||||
serde = { workspace = true }
|
||||
@@ -29,4 +28,3 @@ 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"] }
|
||||
|
||||
@@ -94,7 +94,6 @@ impl AccessControl {
|
||||
}
|
||||
|
||||
/// Get an agent's role.
|
||||
#[allow(dead_code)]
|
||||
pub async fn get_role(&self, agent: &str) -> Option<AgentRole> {
|
||||
self.roles.read().await.get(agent).cloned()
|
||||
}
|
||||
@@ -114,7 +113,6 @@ impl AccessControl {
|
||||
}
|
||||
|
||||
/// Determine which fields should be masked for an agent.
|
||||
#[allow(dead_code)]
|
||||
pub async fn masked_fields(
|
||||
&self,
|
||||
agent: &str,
|
||||
@@ -140,7 +138,6 @@ impl AccessControl {
|
||||
}
|
||||
|
||||
/// Log a query for audit.
|
||||
#[allow(dead_code)]
|
||||
pub async fn log_query(&self, audit: QueryAudit) {
|
||||
self.audit_log.write().await.push(audit);
|
||||
}
|
||||
@@ -152,7 +149,6 @@ impl AccessControl {
|
||||
log[start..].iter().rev().cloned().collect()
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn is_enabled(&self) -> bool {
|
||||
self.enabled
|
||||
}
|
||||
|
||||
+15
-36
@@ -5,51 +5,30 @@ use axum::{
|
||||
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.
|
||||
/// API key auth middleware. Checks X-API-Key header against configured key.
|
||||
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);
|
||||
}
|
||||
// Get the expected key from the request extensions (set by the layer)
|
||||
let expected_key = request.extensions().get::<ApiKey>().cloned();
|
||||
|
||||
let provided = request
|
||||
.headers()
|
||||
.get("x-api-key")
|
||||
.and_then(|v| v.to_str().ok());
|
||||
if let Some(expected) = expected_key {
|
||||
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)
|
||||
match provided {
|
||||
Some(key) if key == expected.0 => {}
|
||||
_ => {
|
||||
tracing::warn!("unauthorized request: missing or invalid API key");
|
||||
return 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
|
||||
Ok(next.run(request).await)
|
||||
}
|
||||
|
||||
/// Wrapper type for the API key, stored in request extensions.
|
||||
|
||||
@@ -1,388 +0,0 @@
|
||||
//! 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
+20
-162
@@ -1,10 +1,8 @@
|
||||
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;
|
||||
@@ -26,30 +24,8 @@ async fn main() {
|
||||
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();
|
||||
// Storage backend
|
||||
let store = storaged::backend::init_local(&config.storage.root);
|
||||
|
||||
// Catalog
|
||||
let registry = catalogd::registry::Registry::new(store.clone());
|
||||
@@ -59,7 +35,7 @@ async fn main() {
|
||||
|
||||
// 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);
|
||||
let engine = queryd::context::QueryEngine::new(registry.clone(), store.clone(), cache);
|
||||
|
||||
// Event journal — append-only mutation log (flush every 100 events)
|
||||
let journal = journald::journal::Journal::new(store.clone(), 100);
|
||||
@@ -77,108 +53,27 @@ async fn main() {
|
||||
// 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;
|
||||
|
||||
// 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("/storage", storaged::service::router(store.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,
|
||||
embed_semaphore: std::sync::Arc::new(tokio::sync::Semaphore::new(1)),
|
||||
.nest("/vectors", vectord::service::router({
|
||||
let index_reg = vectord::index_registry::IndexRegistry::new(store.clone());
|
||||
let _ = index_reg.rebuild().await;
|
||||
vectord::service::VectorState {
|
||||
store: store.clone(),
|
||||
ai_client: ai_client.clone(),
|
||||
job_tracker: vectord::jobs::JobTracker::new(),
|
||||
index_registry: index_reg,
|
||||
hnsw_store: vectord::hnsw::HnswStore::new(),
|
||||
}
|
||||
}))
|
||||
.nest("/workspaces", queryd::workspace_service::router(workspace_mgr))
|
||||
.nest("/journal", journald::service::router(journal))
|
||||
@@ -190,60 +85,23 @@ async fn main() {
|
||||
registry: tool_reg,
|
||||
query_fn: tools::QueryExecutor::new(engine.clone()),
|
||||
}
|
||||
}))
|
||||
// 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
|
||||
},
|
||||
// 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).
|
||||
// Auth middleware (if enabled)
|
||||
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");
|
||||
tracing::info!("API key auth enabled");
|
||||
let api_key = auth::ApiKey(key.clone());
|
||||
app = app.layer(axum::middleware::from_fn_with_state(
|
||||
api_key,
|
||||
auth::api_key_auth,
|
||||
));
|
||||
app = app.layer(axum::Extension(api_key));
|
||||
// Note: auth middleware applied per-route in production
|
||||
// For now, the ApiKey extension is available for handlers to check
|
||||
} 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(DefaultBodyLimit::max(256 * 1024 * 1024)) // 256MB
|
||||
.layer(CorsLayer::new()
|
||||
.allow_origin(Any)
|
||||
.allow_methods(Any)
|
||||
|
||||
@@ -67,7 +67,6 @@ pub struct ToolRegistry {
|
||||
}
|
||||
|
||||
impl ToolRegistry {
|
||||
#[allow(dead_code)]
|
||||
pub fn new() -> Self {
|
||||
let registry = Self {
|
||||
tools: Arc::new(RwLock::new(HashMap::new())),
|
||||
@@ -104,7 +103,7 @@ impl ToolRegistry {
|
||||
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(),
|
||||
sql_template: "SELECT candidate_id, first_name, last_name, phone, email, city, state, zip, vertical, skills, years_experience 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 {
|
||||
|
||||
@@ -7,7 +7,7 @@ use axum::{
|
||||
};
|
||||
use serde::Deserialize;
|
||||
|
||||
use super::registry::{ToolInvocation, ToolRegistry};
|
||||
use super::registry::{Permission, ToolInvocation, ToolRegistry};
|
||||
use crate::tools::ToolState;
|
||||
|
||||
pub fn router(state: ToolState) -> Router {
|
||||
|
||||
@@ -1,226 +0,0 @@
|
||||
//! 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\":["));
|
||||
}
|
||||
}
|
||||
@@ -1,281 +0,0 @@
|
||||
//! 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 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>,
|
||||
/// 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)]
|
||||
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 --
|
||||
|
||||
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");
|
||||
}
|
||||
|
||||
let provider = req.provider.as_deref().unwrap_or("ollama").to_ascii_lowercase();
|
||||
let start_time = chrono::Utc::now();
|
||||
let start_instant = std::time::Instant::now();
|
||||
|
||||
let (resp, used_provider) = match provider.as_str() {
|
||||
"ollama" | "local" | "" => {
|
||||
let r = ollama::chat(&state.ai_client, &req)
|
||||
.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)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::BAD_GATEWAY, format!("ollama cloud: {e}")))?;
|
||||
(r, "ollama_cloud".to_string())
|
||||
}
|
||||
other => {
|
||||
return Err((
|
||||
StatusCode::BAD_REQUEST,
|
||||
format!("unknown provider '{other}' — supported: ollama, ollama_cloud"),
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
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);
|
||||
}
|
||||
}
|
||||
@@ -1,160 +0,0 @@
|
||||
//! 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);
|
||||
}
|
||||
}
|
||||
@@ -1,189 +0,0 @@
|
||||
//! 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.\""));
|
||||
}
|
||||
}
|
||||
@@ -1,150 +0,0 @@
|
||||
//! `/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);
|
||||
}
|
||||
}
|
||||
@@ -1,49 +0,0 @@
|
||||
use axum::{routing::get, Router};
|
||||
use serde::Serialize;
|
||||
use truth::default_truth_store;
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn truth_router() -> Router {
|
||||
Router::new()
|
||||
.route("/context", get(context))
|
||||
}
|
||||
|
||||
#[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,
|
||||
})
|
||||
}
|
||||
@@ -7,8 +7,6 @@ edition = "2024"
|
||||
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 }
|
||||
@@ -23,6 +21,3 @@ csv = { workspace = true }
|
||||
chrono = { workspace = true }
|
||||
object_store = { workspace = true }
|
||||
tokio-postgres = { workspace = true }
|
||||
mysql_async = { workspace = true }
|
||||
uuid = { workspace = true }
|
||||
croner = "3"
|
||||
|
||||
@@ -0,0 +1,235 @@
|
||||
/// ChatGPT PDF conversation parser.
|
||||
/// Extracts structured conversation turns from ChatGPT export PDFs.
|
||||
/// Handles: user/assistant turns, timestamps, conversation titles.
|
||||
|
||||
use arrow::array::{ArrayRef, Int32Array, RecordBatch, StringArray};
|
||||
use arrow::datatypes::{DataType, Field, Schema};
|
||||
use std::sync::Arc;
|
||||
|
||||
/// A single conversation turn.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ConversationTurn {
|
||||
pub source_file: String,
|
||||
pub conversation_title: String,
|
||||
pub turn_number: i32,
|
||||
pub role: String, // "user" or "assistant"
|
||||
pub content: String,
|
||||
pub page_number: i32,
|
||||
}
|
||||
|
||||
/// Parse a ChatGPT PDF into structured conversation turns.
|
||||
pub fn parse_chatgpt_pdf(content: &[u8], filename: &str) -> Result<(Arc<Schema>, Vec<RecordBatch>), String> {
|
||||
let doc = lopdf::Document::load_mem(content)
|
||||
.map_err(|e| format!("PDF load error: {e}"))?;
|
||||
|
||||
// Extract all text
|
||||
let pages = doc.get_pages();
|
||||
let mut all_text = String::new();
|
||||
let mut page_texts: Vec<(i32, 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_texts.push((page_num as i32, text.clone()));
|
||||
all_text.push_str(&text);
|
||||
all_text.push('\n');
|
||||
}
|
||||
}
|
||||
|
||||
if all_text.trim().is_empty() {
|
||||
return Err("PDF contains no extractable text".into());
|
||||
}
|
||||
|
||||
// Try to parse as ChatGPT conversation format
|
||||
let turns = extract_turns(&all_text, filename, &page_texts);
|
||||
|
||||
if turns.is_empty() {
|
||||
// Fall back to page-based extraction if we can't detect conversation structure
|
||||
return parse_as_raw_pages(&page_texts, filename);
|
||||
}
|
||||
|
||||
tracing::info!("parsed ChatGPT PDF '{}': {} turns", filename, turns.len());
|
||||
turns_to_batch(&turns)
|
||||
}
|
||||
|
||||
/// Extract conversation turns from ChatGPT PDF text.
|
||||
/// ChatGPT PDFs typically have patterns like:
|
||||
/// - "You said:" / "ChatGPT said:" (newer format)
|
||||
/// - "User:" / "Assistant:"
|
||||
/// - Or just alternating blocks separated by formatting
|
||||
fn extract_turns(text: &str, filename: &str, page_texts: &[(i32, String)]) -> Vec<ConversationTurn> {
|
||||
let mut turns = Vec::new();
|
||||
let mut turn_num = 0;
|
||||
|
||||
// Try to extract conversation title from first line
|
||||
let title = text.lines().next()
|
||||
.map(|l| l.trim().to_string())
|
||||
.filter(|l| l.len() < 200 && !l.is_empty())
|
||||
.unwrap_or_else(|| filename.replace(".pdf", ""));
|
||||
|
||||
// Strategy 1: Look for "You said:" / "ChatGPT said:" markers (common format)
|
||||
let you_markers = ["You said:", "You:", "User:", "Human:"];
|
||||
let ai_markers = ["ChatGPT said:", "ChatGPT:", "Assistant:", "AI:"];
|
||||
|
||||
let has_markers = you_markers.iter().any(|m| text.contains(m))
|
||||
&& ai_markers.iter().any(|m| text.contains(m));
|
||||
|
||||
if has_markers {
|
||||
let lines: Vec<&str> = text.lines().collect();
|
||||
let mut current_role = String::new();
|
||||
let mut current_content = String::new();
|
||||
let mut current_page = 1i32;
|
||||
|
||||
for line in &lines {
|
||||
let trimmed = line.trim();
|
||||
|
||||
// Check for role markers
|
||||
let mut found_marker = false;
|
||||
for marker in &you_markers {
|
||||
if trimmed.starts_with(marker) || trimmed == *marker {
|
||||
// Save previous turn
|
||||
if !current_content.trim().is_empty() && !current_role.is_empty() {
|
||||
turn_num += 1;
|
||||
turns.push(ConversationTurn {
|
||||
source_file: filename.to_string(),
|
||||
conversation_title: title.clone(),
|
||||
turn_number: turn_num,
|
||||
role: current_role.clone(),
|
||||
content: current_content.trim().to_string(),
|
||||
page_number: current_page,
|
||||
});
|
||||
}
|
||||
current_role = "user".to_string();
|
||||
current_content = trimmed.strip_prefix(marker).unwrap_or("").trim().to_string();
|
||||
current_content.push('\n');
|
||||
found_marker = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if !found_marker {
|
||||
for marker in &ai_markers {
|
||||
if trimmed.starts_with(marker) || trimmed == *marker {
|
||||
if !current_content.trim().is_empty() && !current_role.is_empty() {
|
||||
turn_num += 1;
|
||||
turns.push(ConversationTurn {
|
||||
source_file: filename.to_string(),
|
||||
conversation_title: title.clone(),
|
||||
turn_number: turn_num,
|
||||
role: current_role.clone(),
|
||||
content: current_content.trim().to_string(),
|
||||
page_number: current_page,
|
||||
});
|
||||
}
|
||||
current_role = "assistant".to_string();
|
||||
current_content = trimmed.strip_prefix(marker).unwrap_or("").trim().to_string();
|
||||
current_content.push('\n');
|
||||
found_marker = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !found_marker && !current_role.is_empty() {
|
||||
current_content.push_str(trimmed);
|
||||
current_content.push('\n');
|
||||
}
|
||||
}
|
||||
|
||||
// Save last turn
|
||||
if !current_content.trim().is_empty() && !current_role.is_empty() {
|
||||
turn_num += 1;
|
||||
turns.push(ConversationTurn {
|
||||
source_file: filename.to_string(),
|
||||
conversation_title: title.clone(),
|
||||
turn_number: turn_num,
|
||||
role: current_role,
|
||||
content: current_content.trim().to_string(),
|
||||
page_number: current_page,
|
||||
});
|
||||
}
|
||||
|
||||
return turns;
|
||||
}
|
||||
|
||||
// Strategy 2: Split by paragraph blocks, alternate user/assistant
|
||||
// This is a heuristic fallback for PDFs without clear markers
|
||||
let paragraphs: Vec<&str> = text.split("\n\n")
|
||||
.map(|p| p.trim())
|
||||
.filter(|p| p.len() > 20) // skip short fragments
|
||||
.collect();
|
||||
|
||||
if paragraphs.len() >= 2 {
|
||||
for (i, para) in paragraphs.iter().enumerate() {
|
||||
turn_num += 1;
|
||||
turns.push(ConversationTurn {
|
||||
source_file: filename.to_string(),
|
||||
conversation_title: title.clone(),
|
||||
turn_number: turn_num,
|
||||
role: if i % 2 == 0 { "user" } else { "assistant" }.to_string(),
|
||||
content: para.to_string(),
|
||||
page_number: 1,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
turns
|
||||
}
|
||||
|
||||
/// Fallback: parse as raw pages with full text.
|
||||
fn parse_as_raw_pages(page_texts: &[(i32, String)], filename: &str) -> Result<(Arc<Schema>, Vec<RecordBatch>), String> {
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Field::new("source_file", DataType::Utf8, false),
|
||||
Field::new("conversation_title", DataType::Utf8, false),
|
||||
Field::new("turn_number", DataType::Int32, false),
|
||||
Field::new("role", DataType::Utf8, false),
|
||||
Field::new("content", DataType::Utf8, false),
|
||||
Field::new("page_number", DataType::Int32, false),
|
||||
]));
|
||||
|
||||
let sources: Vec<&str> = page_texts.iter().map(|_| filename).collect();
|
||||
let titles: Vec<String> = page_texts.iter().map(|_| filename.replace(".pdf", "")).collect();
|
||||
let turn_nums: Vec<i32> = (1..=page_texts.len() as i32).collect();
|
||||
let roles: Vec<&str> = page_texts.iter().map(|_| "page").collect();
|
||||
let contents: Vec<&str> = page_texts.iter().map(|(_, t)| t.as_str()).collect();
|
||||
let pages: Vec<i32> = page_texts.iter().map(|(p, _)| *p).collect();
|
||||
|
||||
let arrays: Vec<ArrayRef> = vec![
|
||||
Arc::new(StringArray::from(sources)),
|
||||
Arc::new(StringArray::from(titles.iter().map(|s| s.as_str()).collect::<Vec<_>>())),
|
||||
Arc::new(Int32Array::from(turn_nums)),
|
||||
Arc::new(StringArray::from(roles)),
|
||||
Arc::new(StringArray::from(contents)),
|
||||
Arc::new(Int32Array::from(pages)),
|
||||
];
|
||||
|
||||
let batch = RecordBatch::try_new(schema.clone(), arrays)
|
||||
.map_err(|e| format!("RecordBatch error: {e}"))?;
|
||||
|
||||
Ok((schema, vec![batch]))
|
||||
}
|
||||
|
||||
fn turns_to_batch(turns: &[ConversationTurn]) -> Result<(Arc<Schema>, Vec<RecordBatch>), String> {
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Field::new("source_file", DataType::Utf8, false),
|
||||
Field::new("conversation_title", DataType::Utf8, false),
|
||||
Field::new("turn_number", DataType::Int32, false),
|
||||
Field::new("role", DataType::Utf8, false),
|
||||
Field::new("content", DataType::Utf8, false),
|
||||
Field::new("page_number", DataType::Int32, false),
|
||||
]));
|
||||
|
||||
let arrays: Vec<ArrayRef> = vec![
|
||||
Arc::new(StringArray::from(turns.iter().map(|t| t.source_file.as_str()).collect::<Vec<_>>())),
|
||||
Arc::new(StringArray::from(turns.iter().map(|t| t.conversation_title.as_str()).collect::<Vec<_>>())),
|
||||
Arc::new(Int32Array::from(turns.iter().map(|t| t.turn_number).collect::<Vec<_>>())),
|
||||
Arc::new(StringArray::from(turns.iter().map(|t| t.role.as_str()).collect::<Vec<_>>())),
|
||||
Arc::new(StringArray::from(turns.iter().map(|t| t.content.as_str()).collect::<Vec<_>>())),
|
||||
Arc::new(Int32Array::from(turns.iter().map(|t| t.page_number).collect::<Vec<_>>())),
|
||||
];
|
||||
|
||||
let batch = RecordBatch::try_new(schema.clone(), arrays)
|
||||
.map_err(|e| format!("RecordBatch error: {e}"))?;
|
||||
|
||||
Ok((schema, vec![batch]))
|
||||
}
|
||||
@@ -1,7 +1,5 @@
|
||||
pub mod chatgpt_parser;
|
||||
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;
|
||||
|
||||
@@ -1,410 +0,0 @@
|
||||
/// 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);
|
||||
}
|
||||
}
|
||||
@@ -4,19 +4,8 @@ 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.
|
||||
/// Returns one row per page: (page_number, text_content).
|
||||
/// This handles text-based PDFs. Scanned/image PDFs need OCR (not implemented yet).
|
||||
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}"))?;
|
||||
@@ -37,35 +26,13 @@ pub fn parse_pdf(content: &[u8], source_filename: &str) -> Result<(Arc<Schema>,
|
||||
}
|
||||
}
|
||||
|
||||
// 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,
|
||||
);
|
||||
// PDF has no extractable text — likely scanned/image
|
||||
return Err("PDF contains no extractable text (may be scanned/image — OCR not yet supported)".into());
|
||||
}
|
||||
|
||||
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),
|
||||
@@ -83,197 +50,3 @@ pub fn parse_pdf(content: &[u8], source_filename: &str) -> Result<(Arc<Schema>,
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
@@ -1,308 +0,0 @@
|
||||
/// 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());
|
||||
}
|
||||
}
|
||||
@@ -28,29 +28,12 @@ pub struct IngestResult {
|
||||
}
|
||||
|
||||
/// 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);
|
||||
@@ -60,7 +43,13 @@ pub async fn ingest_file_to_bucket(
|
||||
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::Pdf => {
|
||||
// Try ChatGPT parser first, fall back to generic PDF
|
||||
match crate::chatgpt_parser::parse_chatgpt_pdf(content, filename) {
|
||||
Ok(result) => result,
|
||||
Err(_) => pdf_ingest::parse_pdf(content, filename)?,
|
||||
}
|
||||
}
|
||||
FileType::Text => parse_text_file(content, filename)?,
|
||||
FileType::Unknown => return Err(format!("unknown file type for '{filename}'")),
|
||||
};
|
||||
@@ -115,7 +104,7 @@ pub async fn ingest_file_to_bucket(
|
||||
let schema_fp = fingerprint_schema(&schema);
|
||||
let now = chrono::Utc::now();
|
||||
let obj_ref = ObjectRef {
|
||||
bucket: bucket.to_string(),
|
||||
bucket: "data".to_string(),
|
||||
key: storage_key.clone(),
|
||||
size_bytes: parquet_size,
|
||||
created_at: now,
|
||||
@@ -164,10 +153,6 @@ pub async fn ingest_file_to_bucket(
|
||||
..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),
|
||||
|
||||
@@ -1,725 +0,0 @@
|
||||
//! 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
|
||||
}
|
||||
}
|
||||
@@ -2,11 +2,9 @@
|
||||
/// 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.
|
||||
|
||||
#[allow(unused_imports)]
|
||||
use arrow::datatypes::{DataType, Schema, SchemaRef};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
#[allow(unused_imports)]
|
||||
use std::sync::Arc;
|
||||
|
||||
/// A detected change between two schema versions.
|
||||
|
||||
@@ -1,10 +1,9 @@
|
||||
#[allow(unused_imports)]
|
||||
use axum::{
|
||||
Json, Router,
|
||||
extract::{Multipart, Path, Query, State},
|
||||
http::{HeaderMap, StatusCode},
|
||||
extract::{Multipart, Query, State},
|
||||
http::StatusCode,
|
||||
response::IntoResponse,
|
||||
routing::{delete, get, patch, post},
|
||||
routing::{get, post},
|
||||
};
|
||||
use bytes::Bytes;
|
||||
use object_store::ObjectStore;
|
||||
@@ -12,79 +11,15 @@ use serde::Deserialize;
|
||||
use std::sync::Arc;
|
||||
|
||||
use catalogd::registry::Registry;
|
||||
use crate::{db_ingest, my_stream, pg_stream, pipeline, schedule};
|
||||
use crate::{db_ingest, pipeline};
|
||||
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 {
|
||||
@@ -93,14 +28,6 @@ pub fn router(state: IngestState) -> Router {
|
||||
.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)
|
||||
}
|
||||
|
||||
@@ -116,11 +43,8 @@ struct IngestQuery {
|
||||
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())),
|
||||
@@ -131,34 +55,12 @@ async fn ingest_file(
|
||||
let content = field.bytes().await
|
||||
.map_err(|e| (StatusCode::BAD_REQUEST, format!("read error: {e}")))?;
|
||||
|
||||
tracing::info!(
|
||||
"ingest '{}' ({} bytes) -> bucket={}",
|
||||
filename, content.len(), bucket,
|
||||
);
|
||||
tracing::info!("received file '{}' ({} bytes) for ingest", filename, content.len());
|
||||
|
||||
let dataset_name = query.name.as_deref();
|
||||
|
||||
match pipeline::ingest_file_to_bucket(&filename, &content, dataset_name, &bucket, &store, &state.registry).await {
|
||||
match pipeline::ingest_file(&filename, &content, dataset_name, &state.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 {
|
||||
@@ -286,474 +188,3 @@ async fn import_pg_table(
|
||||
"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");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
/// 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::array::{ArrayRef, RecordBatch, StringArray, UInt64Array};
|
||||
use arrow::datatypes::{DataType, Field, Schema};
|
||||
use chrono::{DateTime, Utc};
|
||||
use object_store::ObjectStore;
|
||||
|
||||
@@ -1,42 +0,0 @@
|
||||
[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"
|
||||
@@ -1,633 +0,0 @@
|
||||
//! 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()),
|
||||
¶ms,
|
||||
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)
|
||||
}
|
||||
|
||||
@@ -15,7 +15,6 @@ 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 }
|
||||
|
||||
+16
-211
@@ -5,7 +5,6 @@ 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;
|
||||
@@ -13,61 +12,27 @@ 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 new(registry: Registry, store: Arc<dyn ObjectStore>, cache: MemCache) -> Self {
|
||||
Self { registry, 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?;
|
||||
@@ -78,16 +43,11 @@ impl QueryEngine {
|
||||
|
||||
/// Pin a dataset into the memory cache.
|
||||
pub async fn pin_dataset(&self, name: &str) -> Result<(), String> {
|
||||
// Read from Parquet
|
||||
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 base_url = Url::parse(&format!("{STORE_SCHEME}://data/"))
|
||||
.map_err(|e| format!("invalid store url: {e}"))?;
|
||||
ctx.runtime_env().register_object_store(&base_url, self.store.clone());
|
||||
|
||||
let dataset = self.registry.get_by_name(name).await
|
||||
.ok_or_else(|| format!("dataset not found: {name}"))?;
|
||||
@@ -98,10 +58,7 @@ impl QueryEngine {
|
||||
|
||||
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()
|
||||
})
|
||||
.filter_map(|o| ListingTableUrl::parse(&format!("{STORE_SCHEME}://data/{}", o.key)).ok())
|
||||
.collect();
|
||||
|
||||
let schema = opts.infer_schema(&ctx.state(), &table_paths[0]).await
|
||||
@@ -122,41 +79,9 @@ impl QueryEngine {
|
||||
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 base_url = Url::parse(&format!("{STORE_SCHEME}://data/"))
|
||||
.map_err(|e| format!("invalid store url: {e}"))?;
|
||||
ctx.runtime_env().register_object_store(&base_url, self.store.clone());
|
||||
|
||||
let datasets = self.registry.list().await;
|
||||
for dataset in &datasets {
|
||||
@@ -164,9 +89,9 @@ impl QueryEngine {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Check cache first — cached batches are bucket-agnostic since
|
||||
// they're already materialized in memory.
|
||||
// Check cache first
|
||||
if let Some(cached) = self.cache.get(&dataset.name).await {
|
||||
// Load any delta files and merge
|
||||
let delta_batches = delta::load_deltas(&self.store, &dataset.name).await.unwrap_or_default();
|
||||
|
||||
let mut all_batches = cached.batches;
|
||||
@@ -181,15 +106,10 @@ impl QueryEngine {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Build ListingTable URLs that include the per-object bucket
|
||||
// scheme. DataFusion routes each scan to the right store
|
||||
// automatically based on the URL.
|
||||
// Fall back to Parquet file
|
||||
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()
|
||||
})
|
||||
.filter_map(|o| ListingTableUrl::parse(&format!("{STORE_SCHEME}://data/{}", o.key)).ok())
|
||||
.collect();
|
||||
|
||||
if table_paths.is_empty() {
|
||||
@@ -204,123 +124,8 @@ impl QueryEngine {
|
||||
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);
|
||||
}
|
||||
}
|
||||
ctx.register_table(&dataset.name, Arc::new(table))
|
||||
.map_err(|e| format!("table registration failed for {}: {e}", dataset.name))?;
|
||||
}
|
||||
|
||||
Ok(ctx)
|
||||
|
||||
+21
-159
@@ -3,15 +3,10 @@
|
||||
/// 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 arrow::array::RecordBatch;
|
||||
use arrow::datatypes::SchemaRef;
|
||||
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};
|
||||
@@ -63,207 +58,74 @@ pub async fn load_deltas(
|
||||
}
|
||||
|
||||
/// 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.
|
||||
/// Optionally deduplicates by a primary key column.
|
||||
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)?;
|
||||
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();
|
||||
|
||||
let has_tombstones = !tombstones.is_empty();
|
||||
let nothing_to_do = delta_batches.is_empty() && !has_tombstones;
|
||||
if nothing_to_do {
|
||||
if delta_batches.is_empty() {
|
||||
return Ok(CompactResult {
|
||||
base_rows: base_batches.iter().map(|b| b.num_rows()).sum(),
|
||||
delta_rows: 0,
|
||||
final_rows: base_batches.iter().map(|b| b.num_rows()).sum(),
|
||||
deltas_merged: 0,
|
||||
tombstones_applied: 0,
|
||||
rows_dropped_by_tombstones: 0,
|
||||
});
|
||||
}
|
||||
|
||||
base_batches.extend(delta_batches);
|
||||
let pre_filter_rows: usize = base_batches.iter().map(|b| b.num_rows()).sum();
|
||||
|
||||
let base_rows: usize = base_batches.iter().map(|b| b.num_rows()).sum();
|
||||
|
||||
// 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.
|
||||
let final_batches = if let Some(_pk) = primary_key_col {
|
||||
// For now, just concatenate. Full dedup requires sorting by PK
|
||||
// and keeping the last row per key — this is a simplification.
|
||||
// TODO: implement proper merge with dedup
|
||||
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}"))?;
|
||||
// Write merged base
|
||||
let mut merged_parquet = Vec::new();
|
||||
for batch in &final_batches {
|
||||
let pq = record_batch_to_parquet(batch)?;
|
||||
merged_parquet.extend_from_slice(&pq);
|
||||
}
|
||||
ops::put(store, base_key, Bytes::from(merged_parquet)).await?;
|
||||
|
||||
// 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.
|
||||
// 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,
|
||||
);
|
||||
tracing::info!("compacted '{}': {} deltas merged, {} → {} rows", dataset_name, delta_count, base_rows, final_rows);
|
||||
|
||||
Ok(CompactResult {
|
||||
base_rows: pre_filter_rows - delta_count, // rough base-before-deltas
|
||||
delta_rows: delta_count,
|
||||
base_rows,
|
||||
delta_rows: final_rows - base_rows + delta_count, // approximate
|
||||
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,
|
||||
}
|
||||
|
||||
@@ -9,6 +9,7 @@ use axum::{
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::cache::CacheStats;
|
||||
use crate::context::QueryEngine;
|
||||
use crate::delta;
|
||||
use crate::paged::ResultStore;
|
||||
@@ -183,31 +184,13 @@ 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))
|
||||
}
|
||||
Ok(result) => Ok(Json(result)),
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,12 +2,15 @@
|
||||
/// Each workspace tracks an agent's activity on a specific contract or search,
|
||||
/// with daily/weekly/monthly tiers and instant handoff capability.
|
||||
|
||||
use arrow::array::{ArrayRef, RecordBatch, StringArray, Int64Array};
|
||||
use arrow::datatypes::{DataType, Field, Schema};
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::delta;
|
||||
use object_store::ObjectStore;
|
||||
|
||||
/// Retention tier for workspace data.
|
||||
|
||||
@@ -15,5 +15,3 @@ bytes = { workspace = true }
|
||||
sha2 = { workspace = true }
|
||||
toml = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
tokio = { workspace = true }
|
||||
async-trait = "0.1"
|
||||
|
||||
@@ -16,46 +16,8 @@ pub struct Config {
|
||||
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")]
|
||||
@@ -66,39 +28,8 @@ pub struct GatewayConfig {
|
||||
|
||||
#[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)]
|
||||
@@ -147,7 +78,6 @@ pub struct ObservabilityConfig {
|
||||
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() }
|
||||
@@ -185,19 +115,13 @@ 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(),
|
||||
},
|
||||
storage: StorageConfig { root: default_storage_root() },
|
||||
catalog: CatalogConfig::default(),
|
||||
query: QueryConfig::default(),
|
||||
sidecar: SidecarConfig { url: default_sidecar_url() },
|
||||
ai: AiConfig::default(),
|
||||
auth: AuthConfig::default(),
|
||||
observability: ObservabilityConfig::default(),
|
||||
agent: AgentSettings::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,4 +3,3 @@ pub mod errors;
|
||||
pub mod arrow_helpers;
|
||||
pub mod config;
|
||||
pub mod pii;
|
||||
pub mod secrets;
|
||||
|
||||
@@ -1,144 +0,0 @@
|
||||
/// 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() }
|
||||
}
|
||||
@@ -116,251 +116,4 @@ pub struct DatasetManifest {
|
||||
/// 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,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -13,4 +13,3 @@ tracing = { workspace = true }
|
||||
object_store = { workspace = true }
|
||||
bytes = { workspace = true }
|
||||
futures = { workspace = true }
|
||||
chrono = { workspace = true }
|
||||
|
||||
@@ -1,330 +0,0 @@
|
||||
/// 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");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,289 +0,0 @@
|
||||
/// 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");
|
||||
}
|
||||
}
|
||||
@@ -1,179 +0,0 @@
|
||||
/// 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(),
|
||||
}
|
||||
}
|
||||
@@ -1,7 +1,3 @@
|
||||
pub mod append_log;
|
||||
pub mod backend;
|
||||
pub mod error_journal;
|
||||
pub mod federation_service;
|
||||
pub mod ops;
|
||||
pub mod registry;
|
||||
pub mod service;
|
||||
|
||||
@@ -1,377 +0,0 @@
|
||||
/// 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)),
|
||||
}
|
||||
}
|
||||
@@ -1,10 +0,0 @@
|
||||
[package]
|
||||
name = "truth"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
tokio = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
@@ -1,523 +0,0 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
|
||||
#[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 },
|
||||
}
|
||||
|
||||
#[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);
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 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,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn default_truth_store() -> TruthStore {
|
||||
let mut store = TruthStore::new();
|
||||
|
||||
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 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");
|
||||
}
|
||||
}
|
||||
@@ -205,13 +205,3 @@ tr:hover td { background: var(--accent-glow); }
|
||||
padding: 8px 12px; border-bottom: 1px solid var(--border); font-size: 13px;
|
||||
}
|
||||
.table-item:hover { background: var(--accent-glow); }
|
||||
|
||||
/* Phase 19 — Playbook panel */
|
||||
.boosted-row { background: rgba(120, 200, 120, 0.10); }
|
||||
.boosted-row td { border-top: 1px solid rgba(120, 200, 120, 0.30); }
|
||||
.mono-cell {
|
||||
font-family: ui-monospace, SFMono-Regular, Menlo, monospace;
|
||||
font-size: 11px; color: var(--text-dim);
|
||||
max-width: 220px; overflow: hidden; text-overflow: ellipsis; white-space: nowrap;
|
||||
}
|
||||
.panel-section + .panel-section { margin-top: 18px; }
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user