use anyhow::Result;
use chrono::{DateTime, Duration, Utc};
use rusqlite::{Connection, params};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SuggestionRow {
pub id: i64,
pub session_id: String,
pub tool_id: String,
pub task_text: Option<String>,
pub score: Option<f64>,
pub suggested_at: String,
pub accepted: bool,
pub accepted_at: Option<String>,
}
pub fn record(
conn: &Connection,
session_id: &str,
tool_id: &str,
task_text: Option<&str>,
score: Option<f64>,
suggested_at: DateTime<Utc>,
) -> Result<i64> {
conn.execute(
"INSERT INTO agent_suggestions
(session_id, tool_id, task_text, score, suggested_at, accepted)
VALUES (?, ?, ?, ?, ?, 0)",
params![
session_id,
tool_id,
task_text,
score,
suggested_at.to_rfc3339(),
],
)?;
Ok(conn.last_insert_rowid())
}
pub fn mark_accepted(
conn: &Connection,
session_id: &str,
tool_id: &str,
accepted_at: DateTime<Utc>,
window_minutes: i64,
) -> Result<usize> {
let cutoff = (accepted_at - Duration::minutes(window_minutes)).to_rfc3339();
let n = conn.execute(
"UPDATE agent_suggestions
SET accepted = 1, accepted_at = ?
WHERE session_id = ?
AND tool_id = ?
AND accepted = 0
AND suggested_at >= ?",
params![accepted_at.to_rfc3339(), session_id, tool_id, cutoff],
)?;
Ok(n)
}
pub fn acceptance_stats(conn: &Connection, since: DateTime<Utc>) -> Result<(i64, i64)> {
let cutoff = since.to_rfc3339();
let row = conn.query_row(
"SELECT
COUNT(*) AS n,
COALESCE(SUM(accepted), 0) AS k
FROM agent_suggestions
WHERE suggested_at >= ?",
params![cutoff],
|r| Ok((r.get::<_, i64>(0)?, r.get::<_, i64>(1)?)),
)?;
Ok(row)
}
pub fn list(conn: &Connection, session_id: Option<&str>) -> Result<Vec<SuggestionRow>> {
let (sql, params): (&str, Vec<Box<dyn rusqlite::ToSql>>) = match session_id {
Some(sid) => (
"SELECT id, session_id, tool_id, task_text, score, suggested_at,
accepted, accepted_at
FROM agent_suggestions
WHERE session_id = ?
ORDER BY suggested_at DESC",
vec![Box::new(sid.to_string())],
),
None => (
"SELECT id, session_id, tool_id, task_text, score, suggested_at,
accepted, accepted_at
FROM agent_suggestions
ORDER BY suggested_at DESC",
vec![],
),
};
let mut stmt = conn.prepare(sql)?;
let param_refs: Vec<&dyn rusqlite::ToSql> = params
.iter()
.map(|b| b.as_ref() as &dyn rusqlite::ToSql)
.collect();
let rows = stmt
.query_map(param_refs.as_slice(), |row| {
Ok(SuggestionRow {
id: row.get(0)?,
session_id: row.get(1)?,
tool_id: row.get(2)?,
task_text: row.get(3)?,
score: row.get(4)?,
suggested_at: row.get(5)?,
accepted: row.get::<_, i64>(6)? != 0,
accepted_at: row.get(7)?,
})
})?
.collect::<Result<Vec<_>, _>>()?;
Ok(rows)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::open;
fn tmp_conn() -> (tempfile::TempDir, Connection) {
let dir = tempfile::tempdir().unwrap();
let conn = open(&dir.path().join("t.sqlite")).unwrap();
(dir, conn)
}
fn seed_tool(conn: &Connection, id: &str) {
conn.execute(
"INSERT OR IGNORE INTO tools (id, type, name, triggers, examples, requires,
enabled, added_at, last_seen_at)
VALUES (?, 'skill', ?, '[]', '[]', '[]', 1,
'2026-05-03T00:00:00+00:00', '2026-05-03T00:00:00+00:00')",
params![id, id],
)
.unwrap();
}
#[test]
fn record_then_mark_accepted_flips_row() {
let (_d, conn) = tmp_conn();
seed_tool(&conn, "skill:caveman");
let suggested = Utc::now();
record(
&conn,
"sess-1",
"skill:caveman",
Some("be terse"),
Some(0.9),
suggested,
)
.unwrap();
let n = mark_accepted(&conn, "sess-1", "skill:caveman", suggested, 60).unwrap();
assert_eq!(n, 1);
let rows = list(&conn, Some("sess-1")).unwrap();
assert_eq!(rows.len(), 1);
assert!(rows[0].accepted);
assert!(rows[0].accepted_at.is_some());
}
#[test]
fn mark_accepted_ignores_old_suggestions() {
let (_d, conn) = tmp_conn();
seed_tool(&conn, "skill:caveman");
let old = Utc::now() - Duration::hours(2);
record(&conn, "sess-1", "skill:caveman", None, None, old).unwrap();
let n = mark_accepted(&conn, "sess-1", "skill:caveman", Utc::now(), 60).unwrap();
assert_eq!(n, 0);
let rows = list(&conn, None).unwrap();
assert!(!rows[0].accepted);
}
#[test]
fn mark_accepted_is_idempotent() {
let (_d, conn) = tmp_conn();
seed_tool(&conn, "skill:x");
let ts = Utc::now();
record(&conn, "s", "skill:x", None, None, ts).unwrap();
assert_eq!(mark_accepted(&conn, "s", "skill:x", ts, 60).unwrap(), 1);
assert_eq!(mark_accepted(&conn, "s", "skill:x", ts, 60).unwrap(), 0);
}
#[test]
fn acceptance_stats_counts_correctly() {
let (_d, conn) = tmp_conn();
seed_tool(&conn, "skill:a");
seed_tool(&conn, "skill:b");
let now = Utc::now();
record(&conn, "s1", "skill:a", None, None, now).unwrap();
record(&conn, "s1", "skill:b", None, None, now).unwrap();
record(&conn, "s2", "skill:a", None, None, now).unwrap();
mark_accepted(&conn, "s1", "skill:a", now, 60).unwrap();
let (n, k) = acceptance_stats(&conn, now - Duration::hours(1)).unwrap();
assert_eq!(n, 3);
assert_eq!(k, 1);
}
}