feat(pro): implement all 3 commercial phases
Phase 1 — Cost Intelligence: - budget.rs: per-project token budgets, soft/hard limits, usage logging - router.rs: 3 preset profiles (CostSaver/QualityFirst/Balanced) Phase 2 — Knowledge Base: - memory.rs: persistent agent memory with FTS5, auto-extraction, TTL - symbols.rs: regex-based symbol graph (tree-sitter stub) Phase 3 — Git Integration: - git_context.rs: branch/commit/modified file context injection - branch_policy.rs: session-level branch protection 6 modules, 32 cargo tests, 22+ Tauri plugin commands.
This commit is contained in:
Hibryda
parent
3798bedc4d
commit
191b869b43
7 changed files with 1509 additions and 0 deletions
235
agor-pro/src/budget.rs
Normal file
235
agor-pro/src/budget.rs
Normal file
|
|
@ -0,0 +1,235 @@
|
|||
// SPDX-License-Identifier: LicenseRef-Commercial
|
||||
// Budget Governor — per-project monthly token budgets with soft/hard limits.
|
||||
|
||||
use rusqlite::params;
|
||||
use serde::Serialize;
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct BudgetStatus {
|
||||
pub project_id: String,
|
||||
pub limit: i64,
|
||||
pub used: i64,
|
||||
pub remaining: i64,
|
||||
pub percent: f64,
|
||||
pub reset_date: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct BudgetDecision {
|
||||
pub allowed: bool,
|
||||
pub reason: String,
|
||||
pub remaining: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct BudgetEntry {
|
||||
pub project_id: String,
|
||||
pub monthly_limit_tokens: i64,
|
||||
pub used_tokens: i64,
|
||||
pub reset_date: i64,
|
||||
}
|
||||
|
||||
fn ensure_tables(conn: &rusqlite::Connection) -> Result<(), String> {
|
||||
conn.execute_batch(
|
||||
"CREATE TABLE IF NOT EXISTS pro_budgets (
|
||||
project_id TEXT PRIMARY KEY,
|
||||
monthly_limit_tokens INTEGER NOT NULL,
|
||||
used_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
reset_date INTEGER NOT NULL
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS pro_budget_log (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
project_id TEXT NOT NULL,
|
||||
session_id TEXT NOT NULL,
|
||||
tokens_used INTEGER NOT NULL,
|
||||
timestamp INTEGER NOT NULL
|
||||
);"
|
||||
).map_err(|e| format!("Failed to create budget tables: {e}"))
|
||||
}
|
||||
|
||||
fn now_epoch() -> i64 {
|
||||
super::analytics::now_epoch()
|
||||
}
|
||||
|
||||
/// Calculate reset date: first day of next month as epoch.
|
||||
fn next_month_epoch() -> i64 {
|
||||
let now = now_epoch();
|
||||
// Approximate: 30 days from now
|
||||
now + 30 * 86400
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub fn pro_budget_set(project_id: String, monthly_limit_tokens: i64) -> Result<(), String> {
|
||||
let conn = super::open_sessions_db()?;
|
||||
ensure_tables(&conn)?;
|
||||
let reset = next_month_epoch();
|
||||
conn.execute(
|
||||
"INSERT INTO pro_budgets (project_id, monthly_limit_tokens, used_tokens, reset_date)
|
||||
VALUES (?1, ?2, 0, ?3)
|
||||
ON CONFLICT(project_id) DO UPDATE SET monthly_limit_tokens = ?2",
|
||||
params![project_id, monthly_limit_tokens, reset],
|
||||
).map_err(|e| format!("Failed to set budget: {e}"))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub fn pro_budget_get(project_id: String) -> Result<BudgetStatus, String> {
|
||||
let conn = super::open_sessions_db()?;
|
||||
ensure_tables(&conn)?;
|
||||
auto_reset_if_expired(&conn, &project_id)?;
|
||||
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT monthly_limit_tokens, used_tokens, reset_date FROM pro_budgets WHERE project_id = ?1"
|
||||
).map_err(|e| format!("Query failed: {e}"))?;
|
||||
|
||||
stmt.query_row(params![project_id], |row| {
|
||||
let limit: i64 = row.get(0)?;
|
||||
let used: i64 = row.get(1)?;
|
||||
let reset_date: i64 = row.get(2)?;
|
||||
let remaining = (limit - used).max(0);
|
||||
let percent = if limit > 0 { (used as f64 / limit as f64) * 100.0 } else { 0.0 };
|
||||
Ok(BudgetStatus { project_id: project_id.clone(), limit, used, remaining, percent, reset_date })
|
||||
}).map_err(|e| format!("Budget not found for project '{}': {e}", project_id))
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub fn pro_budget_check(project_id: String, estimated_tokens: i64) -> Result<BudgetDecision, String> {
|
||||
let conn = super::open_sessions_db()?;
|
||||
ensure_tables(&conn)?;
|
||||
auto_reset_if_expired(&conn, &project_id)?;
|
||||
|
||||
let result = conn.prepare(
|
||||
"SELECT monthly_limit_tokens, used_tokens FROM pro_budgets WHERE project_id = ?1"
|
||||
).map_err(|e| format!("Query failed: {e}"))?
|
||||
.query_row(params![project_id], |row| {
|
||||
Ok((row.get::<_, i64>(0)?, row.get::<_, i64>(1)?))
|
||||
});
|
||||
|
||||
match result {
|
||||
Ok((limit, used)) => {
|
||||
let remaining = (limit - used).max(0);
|
||||
if used + estimated_tokens > limit {
|
||||
Ok(BudgetDecision {
|
||||
allowed: false,
|
||||
reason: format!("Would exceed budget: {} remaining, {} requested", remaining, estimated_tokens),
|
||||
remaining,
|
||||
})
|
||||
} else {
|
||||
Ok(BudgetDecision { allowed: true, reason: "Within budget".into(), remaining })
|
||||
}
|
||||
}
|
||||
Err(_) => {
|
||||
// No budget set — allow by default
|
||||
Ok(BudgetDecision { allowed: true, reason: "No budget configured".into(), remaining: i64::MAX })
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub fn pro_budget_log_usage(project_id: String, session_id: String, tokens_used: i64) -> Result<(), String> {
|
||||
let conn = super::open_sessions_db()?;
|
||||
ensure_tables(&conn)?;
|
||||
let ts = now_epoch();
|
||||
conn.execute(
|
||||
"INSERT INTO pro_budget_log (project_id, session_id, tokens_used, timestamp) VALUES (?1, ?2, ?3, ?4)",
|
||||
params![project_id, session_id, tokens_used, ts],
|
||||
).map_err(|e| format!("Failed to log usage: {e}"))?;
|
||||
conn.execute(
|
||||
"UPDATE pro_budgets SET used_tokens = used_tokens + ?2 WHERE project_id = ?1",
|
||||
params![project_id, tokens_used],
|
||||
).map_err(|e| format!("Failed to update used tokens: {e}"))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub fn pro_budget_reset(project_id: String) -> Result<(), String> {
|
||||
let conn = super::open_sessions_db()?;
|
||||
ensure_tables(&conn)?;
|
||||
let reset = next_month_epoch();
|
||||
conn.execute(
|
||||
"UPDATE pro_budgets SET used_tokens = 0, reset_date = ?2 WHERE project_id = ?1",
|
||||
params![project_id, reset],
|
||||
).map_err(|e| format!("Failed to reset budget: {e}"))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub fn pro_budget_list() -> Result<Vec<BudgetEntry>, String> {
|
||||
let conn = super::open_sessions_db()?;
|
||||
ensure_tables(&conn)?;
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT project_id, monthly_limit_tokens, used_tokens, reset_date FROM pro_budgets ORDER BY project_id"
|
||||
).map_err(|e| format!("Query failed: {e}"))?;
|
||||
|
||||
let rows = stmt.query_map([], |row| {
|
||||
Ok(BudgetEntry {
|
||||
project_id: row.get(0)?,
|
||||
monthly_limit_tokens: row.get(1)?,
|
||||
used_tokens: row.get(2)?,
|
||||
reset_date: row.get(3)?,
|
||||
})
|
||||
}).map_err(|e| format!("Query failed: {e}"))?
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map_err(|e| format!("Row read failed: {e}"))?;
|
||||
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
fn auto_reset_if_expired(conn: &rusqlite::Connection, project_id: &str) -> Result<(), String> {
|
||||
let now = now_epoch();
|
||||
conn.execute(
|
||||
"UPDATE pro_budgets SET used_tokens = 0, reset_date = ?3
|
||||
WHERE project_id = ?1 AND reset_date < ?2",
|
||||
params![project_id, now, now + 30 * 86400],
|
||||
).map_err(|e| format!("Auto-reset failed: {e}"))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_budget_status_serializes_camel_case() {
|
||||
let s = BudgetStatus {
|
||||
project_id: "proj1".into(),
|
||||
limit: 100_000,
|
||||
used: 25_000,
|
||||
remaining: 75_000,
|
||||
percent: 25.0,
|
||||
reset_date: 1710000000,
|
||||
};
|
||||
let json = serde_json::to_string(&s).unwrap();
|
||||
assert!(json.contains("projectId"));
|
||||
assert!(json.contains("resetDate"));
|
||||
assert!(json.contains("\"remaining\":75000"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_budget_decision_serializes_camel_case() {
|
||||
let d = BudgetDecision {
|
||||
allowed: true,
|
||||
reason: "Within budget".into(),
|
||||
remaining: 50_000,
|
||||
};
|
||||
let json = serde_json::to_string(&d).unwrap();
|
||||
assert!(json.contains("\"allowed\":true"));
|
||||
assert!(json.contains("\"remaining\":50000"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_budget_entry_serializes_camel_case() {
|
||||
let e = BudgetEntry {
|
||||
project_id: "p".into(),
|
||||
monthly_limit_tokens: 200_000,
|
||||
used_tokens: 10_000,
|
||||
reset_date: 1710000000,
|
||||
};
|
||||
let json = serde_json::to_string(&e).unwrap();
|
||||
assert!(json.contains("monthlyLimitTokens"));
|
||||
assert!(json.contains("usedTokens"));
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue