use super::super::ToolSpanRecord;
use anyhow::Result;
use rusqlite::{Connection, params};
const CLEAR_PATHS_SQL: &str = "
DELETE FROM tool_span_paths
WHERE span_id IN (SELECT span_id FROM tool_spans WHERE session_id = ?1)";
const CLEAR_SPANS_SQL: &str = "DELETE FROM tool_spans WHERE session_id = ?1";
const DELETE_PATHS_SQL: &str = "DELETE FROM tool_span_paths WHERE span_id = ?1";
const INSERT_PATH_SQL: &str = "INSERT INTO tool_span_paths (span_id, path) VALUES (?1, ?2)";
const UPSERT_SPAN_SQL: &str = "
INSERT INTO tool_spans (
span_id, session_id, tool, tool_call_id, status,
started_at_ms, ended_at_ms, lead_time_ms,
tokens_in, tokens_out, reasoning_tokens, cost_usd_e6, paths_json,
parent_span_id, depth, subtree_cost_usd_e6, subtree_token_count
) VALUES (?1,?2,?3,?4,?5,?6,?7,?8,?9,?10,?11,?12,?13,?14,?15,?16,?17)
ON CONFLICT(span_id) DO UPDATE SET
session_id=excluded.session_id,
tool=excluded.tool,
tool_call_id=excluded.tool_call_id,
status=excluded.status,
started_at_ms=excluded.started_at_ms,
ended_at_ms=excluded.ended_at_ms,
lead_time_ms=excluded.lead_time_ms,
tokens_in=excluded.tokens_in,
tokens_out=excluded.tokens_out,
reasoning_tokens=excluded.reasoning_tokens,
cost_usd_e6=excluded.cost_usd_e6,
paths_json=excluded.paths_json,
parent_span_id=excluded.parent_span_id,
depth=excluded.depth,
subtree_cost_usd_e6=excluded.subtree_cost_usd_e6,
subtree_token_count=excluded.subtree_token_count";
pub(crate) fn clear_session_spans(conn: &Connection, session_id: &str) -> Result<()> {
conn.execute(CLEAR_PATHS_SQL, params![session_id])?;
conn.execute(CLEAR_SPANS_SQL, params![session_id])?;
Ok(())
}
pub(crate) fn upsert_tool_span_record(
conn: &Connection,
span: &ToolSpanRecord,
) -> Result<ToolSpanRecord> {
let persisted = namespaced_record(span);
upsert_span(conn, &persisted)?;
replace_paths(conn, &persisted)?;
Ok(persisted)
}
fn namespaced_record(span: &ToolSpanRecord) -> ToolSpanRecord {
ToolSpanRecord {
span_id: persisted_span_id(&span.session_id, &span.span_id),
parent_span_id: span
.parent_span_id
.as_deref()
.map(|id| persisted_span_id(&span.session_id, id)),
..span.clone()
}
}
fn persisted_span_id(session_id: &str, span_id: &str) -> String {
format!("v1:{}:{session_id}:{span_id}", session_id.len())
}
fn upsert_span(conn: &Connection, span: &ToolSpanRecord) -> Result<()> {
conn.execute(
UPSERT_SPAN_SQL,
params![
&span.span_id,
&span.session_id,
span.tool.as_deref(),
span.tool_call_id.as_deref(),
&span.status,
span.started_at_ms.map(|value| value as i64),
span.ended_at_ms.map(|value| value as i64),
span.lead_time_ms.map(|value| value as i64),
span.tokens_in.map(|value| value as i64),
span.tokens_out.map(|value| value as i64),
span.reasoning_tokens.map(|value| value as i64),
span.cost_usd_e6,
serde_json::to_string(&span.paths)?,
span.parent_span_id.as_deref(),
span.depth as i64,
span.subtree_cost_usd_e6,
span.subtree_token_count.map(|value| value as i64),
],
)?;
Ok(())
}
fn replace_paths(conn: &Connection, span: &ToolSpanRecord) -> Result<()> {
conn.execute(DELETE_PATHS_SQL, params![&span.span_id])?;
for path in &span.paths {
conn.execute(INSERT_PATH_SQL, params![&span.span_id, path])?;
}
Ok(())
}