use khive_runtime::RuntimeError;
use khive_storage::types::{SqlStatement, SqlValue};
use khive_storage::SqlAccess;
pub use khive_brain_core::compute_query_class;
fn sql_err(context: &str, e: impl std::fmt::Display) -> RuntimeError {
RuntimeError::Internal(format!("serve ledger {context}: {e}"))
}
pub struct ServeLedgerRow {
pub id: String,
#[allow(dead_code)]
pub namespace: String,
#[allow(dead_code)]
pub consumer_kind: String,
pub served_by_profile_id: Option<String>,
pub resolved_profile_id: Option<String>,
#[allow(dead_code)]
pub resolved_at: Option<i64>,
pub accounting_profile_id: Option<String>,
#[allow(dead_code)]
pub target_id: String,
#[allow(dead_code)]
pub query_class: String,
#[allow(dead_code)]
pub query_raw: String,
#[allow(dead_code)]
pub served_at: i64,
#[allow(dead_code)]
pub grade: Option<String>,
#[allow(dead_code)]
pub graded_at: Option<i64>,
pub scorer_run_id: Option<String>,
}
fn row_text(row: &khive_storage::types::SqlRow, col: &str) -> Result<String, RuntimeError> {
match row.get(col) {
Some(SqlValue::Text(s)) => Ok(s.clone()),
_ => Err(sql_err("read row", format!("missing {col} column"))),
}
}
fn row_opt_text(row: &khive_storage::types::SqlRow, col: &str) -> Option<String> {
match row.get(col) {
Some(SqlValue::Text(s)) => Some(s.clone()),
_ => None,
}
}
fn row_opt_int(row: &khive_storage::types::SqlRow, col: &str) -> Option<i64> {
match row.get(col) {
Some(SqlValue::Integer(n)) => Some(*n),
_ => None,
}
}
fn row_int(row: &khive_storage::types::SqlRow, col: &str) -> Result<i64, RuntimeError> {
match row.get(col) {
Some(SqlValue::Integer(n)) => Ok(*n),
_ => Err(sql_err("read row", format!("missing {col} column"))),
}
}
#[allow(clippy::too_many_arguments)]
pub async fn record_serve(
sql: &dyn SqlAccess,
id: &str,
namespace: &str,
consumer_kind: &str,
served_by_profile_id: Option<&str>,
resolved_profile_id: Option<&str>,
resolved_at: Option<i64>,
target_id: &str,
query_class: &str,
query_raw: &str,
served_at: i64,
) -> Result<bool, RuntimeError> {
let mut writer = sql.writer().await.map_err(|e| sql_err("writer", e))?;
let result = writer
.execute(SqlStatement {
sql: "INSERT INTO brain_serve_ledger \
(id, namespace, consumer_kind, served_by_profile_id, resolved_profile_id, \
resolved_at, target_id, query_class, query_raw, served_at) \
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)"
.into(),
params: vec![
SqlValue::Text(id.to_string()),
SqlValue::Text(namespace.to_string()),
SqlValue::Text(consumer_kind.to_string()),
served_by_profile_id
.map(|s| SqlValue::Text(s.to_string()))
.unwrap_or(SqlValue::Null),
resolved_profile_id
.map(|s| SqlValue::Text(s.to_string()))
.unwrap_or(SqlValue::Null),
resolved_at.map(SqlValue::Integer).unwrap_or(SqlValue::Null),
SqlValue::Text(target_id.to_string()),
SqlValue::Text(query_class.to_string()),
SqlValue::Text(query_raw.to_string()),
SqlValue::Integer(served_at),
],
label: Some("brain_serve_ledger_insert".into()),
})
.await;
match result {
Ok(_) => Ok(true),
Err(e) if e.is_unique_constraint_violation() => {
if natural_key_row_exists(sql, namespace, target_id, query_class, served_at).await? {
Ok(false)
} else {
Err(sql_err("insert serve row", e))
}
}
Err(e) => Err(sql_err("insert serve row", e)),
}
}
async fn natural_key_row_exists(
sql: &dyn SqlAccess,
namespace: &str,
target_id: &str,
query_class: &str,
served_at: i64,
) -> Result<bool, RuntimeError> {
let mut reader = sql.reader().await.map_err(|e| sql_err("reader", e))?;
let row = reader
.query_row(SqlStatement {
sql: "SELECT id FROM brain_serve_ledger \
WHERE namespace = ?1 AND target_id = ?2 AND query_class = ?3 AND served_at = ?4"
.into(),
params: vec![
SqlValue::Text(namespace.to_string()),
SqlValue::Text(target_id.to_string()),
SqlValue::Text(query_class.to_string()),
SqlValue::Integer(served_at),
],
label: Some("brain_serve_ledger_natural_key_check".into()),
})
.await
.map_err(|e| sql_err("check natural key", e))?;
Ok(row.is_some())
}
pub async fn get_serve_row(
sql: &dyn SqlAccess,
id: &str,
) -> Result<Option<ServeLedgerRow>, RuntimeError> {
let mut reader = sql.reader().await.map_err(|e| sql_err("reader", e))?;
let row = reader
.query_row(SqlStatement {
sql: "SELECT id, namespace, consumer_kind, served_by_profile_id, \
resolved_profile_id, resolved_at, accounting_profile_id, target_id, \
query_class, query_raw, served_at, grade, graded_at, scorer_run_id \
FROM brain_serve_ledger WHERE id = ?1"
.into(),
params: vec![SqlValue::Text(id.to_string())],
label: Some("brain_serve_ledger_get".into()),
})
.await
.map_err(|e| sql_err("read row", e))?;
let Some(row) = row else {
return Ok(None);
};
Ok(Some(ServeLedgerRow {
id: row_text(&row, "id")?,
namespace: row_text(&row, "namespace")?,
consumer_kind: row_text(&row, "consumer_kind")?,
served_by_profile_id: row_opt_text(&row, "served_by_profile_id"),
resolved_profile_id: row_opt_text(&row, "resolved_profile_id"),
resolved_at: row_opt_int(&row, "resolved_at"),
accounting_profile_id: row_opt_text(&row, "accounting_profile_id"),
target_id: row_text(&row, "target_id")?,
query_class: row_text(&row, "query_class")?,
query_raw: row_text(&row, "query_raw")?,
served_at: row_int(&row, "served_at")?,
grade: row_opt_text(&row, "grade"),
graded_at: row_opt_int(&row, "graded_at"),
scorer_run_id: row_opt_text(&row, "scorer_run_id"),
}))
}
pub async fn backfill_grade(
sql: &dyn SqlAccess,
id: &str,
grade: &str,
graded_at_us: i64,
scorer_run_id: &str,
) -> Result<(), RuntimeError> {
let mut writer = sql.writer().await.map_err(|e| sql_err("writer", e))?;
writer
.execute(SqlStatement {
sql: "UPDATE brain_serve_ledger SET grade = ?1, graded_at = ?2, scorer_run_id = ?3 \
WHERE id = ?4"
.into(),
params: vec![
SqlValue::Text(grade.to_string()),
SqlValue::Integer(graded_at_us),
SqlValue::Text(scorer_run_id.to_string()),
SqlValue::Text(id.to_string()),
],
label: Some("brain_serve_ledger_backfill_grade".into()),
})
.await
.map_err(|e| sql_err("backfill grade", e))?;
Ok(())
}
pub enum ServeLedgerResolution {
AlreadyGraded,
Proceed {
accounting_profile_id: Option<String>,
},
NotFound,
}
pub async fn resolve(
sql: &dyn SqlAccess,
serve_ledger_id: &str,
scorer_run_id: &str,
) -> Result<ServeLedgerResolution, RuntimeError> {
let Some(row) = get_serve_row(sql, serve_ledger_id).await? else {
return Ok(ServeLedgerResolution::NotFound);
};
if row.scorer_run_id.as_deref() == Some(scorer_run_id) {
return Ok(ServeLedgerResolution::AlreadyGraded);
}
Ok(ServeLedgerResolution::Proceed {
accounting_profile_id: row.accounting_profile_id,
})
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn record_serve_routes_through_writer_task_when_flag_enabled() {
let dir = tempfile::tempdir().expect("tempdir");
let db_path = dir.path().join("serve-ledger-write-queue-routing.db");
let pool_cfg = khive_db::PoolConfig {
path: Some(db_path),
write_queue_enabled: true,
..khive_db::PoolConfig::default()
};
let pool = std::sync::Arc::new(khive_db::ConnectionPool::new(pool_cfg).expect("pool"));
{
let mut writer = pool.writer().expect("writer");
khive_db::run_migrations(writer.conn_mut()).expect("migrations");
}
let sql: std::sync::Arc<dyn SqlAccess> =
std::sync::Arc::new(khive_db::SqlBridge::new(std::sync::Arc::clone(&pool), true));
let writer_task = pool
.writer_task_handle()
.unwrap()
.expect("writer task must be spawned with the flag on for a file-backed pool");
let (started_tx, started_rx) = tokio::sync::oneshot::channel::<()>();
let (release_tx, release_rx) = tokio::sync::oneshot::channel::<()>();
let occupier = {
let writer_task = writer_task.clone();
tokio::spawn(async move {
writer_task
.send(move |_conn| {
let _ = started_tx.send(());
let _ = release_rx.blocking_recv();
Ok::<(), khive_storage::StorageError>(())
})
.await
})
};
started_rx
.await
.expect("occupier must signal it has started running inside the writer task");
assert_eq!(
writer_task.queue_depth(),
0,
"channel must start empty once the occupier has been dequeued and is running"
);
let record_task = tokio::spawn(async move {
record_serve(
sql.as_ref(),
"write-queue-routing-row",
"local",
"recall",
Some("profile-a"),
None,
None,
"write-queue-routing-target",
"write-queue-routing-class",
"write queue routing test query",
1_700_000_000_000_000,
)
.await
});
let mut saw_enqueued = false;
for _ in 0..100 {
if writer_task.queue_depth() >= 1 {
saw_enqueued = true;
break;
}
tokio::time::sleep(std::time::Duration::from_millis(5)).await;
}
assert!(
saw_enqueued,
"record_serve's write request never appeared in the writer task's \
channel while the occupier held the single drain slot — \
SqliteWriter::execute is not routing this single-row write \
through the shared writer task"
);
release_tx
.send(())
.expect("occupier must still be waiting on the release signal");
occupier
.await
.expect("occupier task must not panic")
.expect("occupier write must succeed");
let inserted = record_task
.await
.expect("record_task must not panic")
.expect("record_serve must succeed once unblocked");
assert!(inserted, "record_serve must report a fresh row insert");
}
}