use anyhow::Result;
use chrono::Utc;
use rusqlite::params;
use serde_json::Value;
use sqlitegraph::{GraphEdge, GraphEntity, SqliteGraph};
pub mod audit;
pub mod discovery;
pub mod evidence;
pub mod handoff;
pub mod knowledge;
pub mod magellan_bridge;
pub mod navigation;
pub mod ontology;
pub mod planning;
pub mod search;
pub mod types;
pub mod wiki;
pub use planning::{KanbanStatus, KanbanUpdate};
pub use types::{
ActionRecord, ActionTrace, AppliedKanbanUpdate, AtheneumError, BlockerType, CommitParams,
EdgeType, EndSessionParams, EntityType, FileWriteParams, FixChainParams, GraphStats, Neighbors,
OntologyClassInfo, OntologyPropertyInfo, PromptParams, RecordEventParams, RequirementStatus,
SearchResult, SessionParams, SessionSummary, SubgraphView, TaskDetail, TestRunParams,
ToolCallParams, ToolCallRecord, ToolCallTrace, ONTOLOGY_CLASS_KIND, ONTOLOGY_PROPERTY_KIND,
};
pub use wiki::{
content_hash, extract_kanban_updates, extract_wikilinks, parse_journal_sections,
JournalSection, WikiPage,
};
pub(super) fn json_to_string(v: &Value) -> Result<String> {
serde_json::to_string(v).map_err(|e| anyhow::anyhow!("JSON serialization failed: {}", e))
}
pub struct AtheneumGraph {
inner: SqliteGraph,
}
impl AtheneumGraph {
pub fn open_in_memory() -> Result<Self> {
let inner = SqliteGraph::open_in_memory()?;
let g = Self { inner };
g.run_startup_migrations()?;
Ok(g)
}
pub fn open(path: &std::path::Path) -> Result<Self> {
let inner = SqliteGraph::open(path)?;
let g = Self { inner };
g.run_startup_migrations()?;
Ok(g)
}
fn run_startup_migrations(&self) -> Result<()> {
self.with_raw_connection(crate::db::run_migrations)
}
pub fn with_raw_connection<F, R>(&self, f: F) -> Result<R>
where
F: FnOnce(&rusqlite::Connection) -> Result<R>,
{
if let Some(direct) = self.inner.pool.direct_connection() {
f(direct)
} else {
let pooled = self
.inner
.pool
.get()
.map_err(|e| anyhow::anyhow!("Failed to get connection: {}", e))?;
f(&pooled)
}
}
pub fn is_healthy(&self) -> bool {
self.with_raw_connection(|conn| {
conn.execute_batch("SELECT 1").ok();
Ok(())
})
.is_ok()
}
pub fn get_entity(&self, id: i64) -> Result<GraphEntity> {
self.inner
.get_entity(id)
.map_err(|_e| AtheneumError::EntityNotFound(id).into())
}
pub fn get_edge(&self, id: i64) -> Result<GraphEdge> {
self.inner
.get_edge(id)
.map_err(|_e| AtheneumError::EdgeNotFound(id).into())
}
pub fn outgoing_edges(&self, entity_id: i64) -> Result<Vec<GraphEdge>> {
self.with_raw_connection(|conn| {
let mut stmt = conn.prepare_cached(
"SELECT id, from_id, to_id, edge_type, data FROM graph_edges WHERE from_id=?1 ORDER BY id"
)?;
let rows = stmt.query_map(rusqlite::params![entity_id], |row| {
Ok(GraphEdge {
id: row.get(0)?,
from_id: row.get(1)?,
to_id: row.get(2)?,
edge_type: row.get(3)?,
data: serde_json::from_str(row.get_ref(4)?.as_str()?)
.map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?,
})
})?;
let mut edges = Vec::new();
for row in rows {
edges.push(row?);
}
Ok(edges)
})
}
pub fn incoming_edges(&self, entity_id: i64) -> Result<Vec<GraphEdge>> {
self.with_raw_connection(|conn| {
let mut stmt = conn.prepare_cached(
"SELECT id, from_id, to_id, edge_type, data FROM graph_edges WHERE to_id=?1 ORDER BY id"
)?;
let rows = stmt.query_map(rusqlite::params![entity_id], |row| {
Ok(GraphEdge {
id: row.get(0)?,
from_id: row.get(1)?,
to_id: row.get(2)?,
edge_type: row.get(3)?,
data: serde_json::from_str(row.get_ref(4)?.as_str()?)
.map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?,
})
})?;
let mut edges = Vec::new();
for row in rows {
edges.push(row?);
}
Ok(edges)
})
}
pub fn entities_by_kind(&self, kind: &str) -> Result<Vec<GraphEntity>> {
with_graph_conn(&self.inner, |conn| {
let mut stmt = conn.prepare_cached(
"SELECT id, kind, name, file_path, data FROM graph_entities WHERE kind=?1",
)?;
let rows = stmt.query_map(params![kind], |row| {
Ok(GraphEntity {
id: row.get(0)?,
kind: row.get(1)?,
name: row.get(2)?,
file_path: row.get(3)?,
data: serde_json::from_str(row.get_ref(4)?.as_str()?)
.map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?,
})
})?;
let mut entities = Vec::new();
for row in rows {
entities.push(row?);
}
Ok(entities)
})
}
pub fn count_entities_by_kind(&self) -> Result<Vec<(String, i64)>> {
with_graph_conn(&self.inner, |conn| {
let mut stmt = conn.prepare_cached(
"SELECT kind, COUNT(*) as count FROM graph_entities GROUP BY kind ORDER BY count DESC"
)?;
let rows = stmt.query_map([], |row| {
Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)?))
})?;
let mut counts = Vec::new();
for row in rows {
counts.push(row?);
}
Ok(counts)
})
}
pub fn count_edges_by_type(&self) -> Result<Vec<(String, i64)>> {
with_graph_conn(&self.inner, |conn| {
let mut stmt = conn.prepare_cached(
"SELECT edge_type, COUNT(*) as count FROM graph_edges GROUP BY edge_type ORDER BY count DESC"
)?;
let rows = stmt.query_map([], |row| {
Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)?))
})?;
let mut counts = Vec::new();
for row in rows {
counts.push(row?);
}
Ok(counts)
})
}
pub fn insert_agent(&self, name: &str, data: Value) -> Result<i64> {
let metadata_str = json_to_string(&data)?;
let sql_id = self.with_raw_connection(|conn| {
let project_id = data.get("project_id").and_then(|v| v.as_str());
conn.execute(
"INSERT OR IGNORE INTO agents (name, project_id, metadata, created_at)
VALUES (?1, ?2, ?3, ?4)",
rusqlite::params![name, project_id, metadata_str, Utc::now().to_rfc3339()],
)?;
let id: i64 = conn.query_row(
"SELECT id FROM agents WHERE name = ?1",
rusqlite::params![name],
|row| row.get(0),
)?;
Ok(id)
})?;
let mut data = data;
if let Some(obj) = data.as_object_mut() {
obj.insert("sql_id".to_string(), Value::Number(sql_id.into()));
}
let entity = GraphEntity {
id: 0,
kind: EntityType::Agent.as_str().to_string(),
name: name.to_string(),
file_path: None,
data,
};
self.inner.insert_entity(&entity).map_err(Into::into)
}
pub fn insert_task(&self, name: &str, data: Value) -> Result<i64> {
let entity = GraphEntity {
id: 0,
kind: EntityType::Task.as_str().to_string(),
name: name.to_string(),
file_path: None,
data,
};
self.inner.insert_entity(&entity).map_err(Into::into)
}
pub fn insert_event(&self, name: &str, mut data: Value) -> Result<i64> {
if let Some(obj) = data.as_object_mut() {
if !obj.contains_key("timestamp") {
obj.insert(
"timestamp".to_string(),
Value::String(Utc::now().to_rfc3339()),
);
}
}
let entity = GraphEntity {
id: 0,
kind: EntityType::Event.as_str().to_string(),
name: name.to_string(),
file_path: None,
data,
};
self.inner.insert_entity(&entity).map_err(Into::into)
}
pub fn insert_edge(
&self,
from_id: i64,
to_id: i64,
edge_type: EdgeType,
data: Value,
) -> Result<i64> {
let edge = GraphEdge {
id: 0,
from_id,
to_id,
edge_type: edge_type.as_str().to_string(),
data,
};
self.inner.insert_edge(&edge).map_err(Into::into)
}
pub fn events_performed_by(&self, agent_id: i64) -> Result<Vec<GraphEntity>> {
let mut events = Vec::new();
let edges = get_incoming_edges(&self.inner, agent_id)?
.into_iter()
.filter(|e| e.edge_type == EdgeType::PerformedBy.as_str());
for edge in edges {
if let Ok(entity) = self.get_entity(edge.from_id) {
if entity.kind == EntityType::Event.as_str() {
events.push(entity);
}
}
}
Ok(events)
}
pub fn tasks_assigned_to(&self, agent_id: i64) -> Result<Vec<GraphEntity>> {
let mut tasks = Vec::new();
let edges = get_incoming_edges(&self.inner, agent_id)?
.into_iter()
.filter(|e| e.edge_type == EdgeType::AssignedTo.as_str());
for edge in edges {
if let Ok(entity) = self.get_entity(edge.from_id) {
if entity.kind == EntityType::Task.as_str() {
tasks.push(entity);
}
}
}
Ok(tasks)
}
pub fn causal_chain(&self, event_id: i64) -> Result<Vec<GraphEntity>> {
let mut chain = Vec::new();
let mut current = Some(event_id);
let mut visited = std::collections::HashSet::new();
while let Some(id) = current {
if !visited.insert(id) {
break;
}
if let Ok(entity) = self.get_entity(id) {
chain.push(entity);
if let Ok(edges) = get_outgoing_edges(&self.inner, id) {
current = edges
.into_iter()
.find(|e| e.edge_type == EdgeType::CausedBy.as_str())
.map(|e| e.to_id);
} else {
current = None;
}
} else {
current = None;
}
}
Ok(chain)
}
pub(super) fn find_entity_id_by_data(
&self,
kind: &str,
key: &str,
value: &str,
) -> Result<Option<i64>> {
self.with_raw_connection(|conn| {
Ok(conn
.query_row(
"SELECT id FROM graph_entities WHERE kind = ?1 AND json_extract(data, ?2) = ?3 LIMIT 1",
rusqlite::params![kind, format!("$.{}", key), value],
|r| r.get(0),
)
.ok())
})
}
pub(super) fn find_entity_id_by_kind_and_name(
&self,
kind: &str,
name: &str,
) -> Result<Option<i64>> {
self.with_raw_connection(|conn| {
Ok(conn
.query_row(
"SELECT id FROM graph_entities WHERE kind = ?1 AND name = ?2 LIMIT 1",
rusqlite::params![kind, name],
|r| r.get(0),
)
.ok())
})
}
fn update_entity_data(&self, id: i64, data: &Value) -> Result<()> {
with_graph_conn(&self.inner, |conn| {
conn.execute(
"UPDATE graph_entities SET data = ?1 WHERE id = ?2",
params![serde_json::to_string(data)?, id],
)?;
Ok(())
})
}
}
fn parse_frontmatter(content: &str) -> Result<(Value, &str)> {
let first_marker = content
.find("---")
.ok_or_else(|| anyhow::anyhow!("No frontmatter start marker"))?;
let second_marker = content[first_marker + 3..]
.find("---")
.ok_or_else(|| anyhow::anyhow!("No frontmatter end marker"))?
+ first_marker
+ 3;
let frontmatter_text = &content[first_marker + 3..second_marker];
let mut map = serde_json::Map::new();
for line in frontmatter_text.lines() {
let line = line.trim();
if line.is_empty() || line.starts_with('#') {
continue;
}
let Some((key, rest)) = line.split_once(':') else {
continue;
};
let key = key.trim();
let rest = rest.trim();
let value = if rest.starts_with('[') && rest.ends_with(']') {
let inner = &rest[1..rest.len() - 1];
let items: Vec<Value> = inner
.split(',')
.map(|s| Value::String(s.trim().trim_matches('"').trim_matches('\'').to_string()))
.collect();
Value::Array(items)
} else if rest == "true" {
Value::Bool(true)
} else if rest == "false" {
Value::Bool(false)
} else if let Ok(n) = rest.parse::<i64>() {
Value::Number(n.into())
} else if let Ok(n) = rest.parse::<f64>() {
serde_json::Number::from_f64(n)
.map(Value::Number)
.unwrap_or_else(|| Value::String(rest.to_string()))
} else {
Value::String(rest.trim_matches('"').trim_matches('\'').to_string())
};
map.insert(key.to_string(), value);
}
let body = &content[second_marker + 3..];
Ok((Value::Object(map), body))
}
pub(super) fn with_graph_conn<F, R>(graph: &SqliteGraph, f: F) -> Result<R>
where
F: FnOnce(&rusqlite::Connection) -> Result<R>,
{
if let Some(direct) = graph.pool.direct_connection() {
f(direct)
} else {
let pooled = graph
.pool
.get()
.map_err(|e| anyhow::anyhow!("Failed to get connection: {}", e))?;
f(&pooled)
}
}
fn get_incoming_edges(graph: &SqliteGraph, entity_id: i64) -> Result<Vec<GraphEdge>> {
with_graph_conn(graph, |conn| {
let mut stmt = conn.prepare_cached(
"SELECT id, from_id, to_id, edge_type, data FROM graph_edges WHERE to_id=?1 ORDER BY id",
)?;
let rows = stmt.query_map(params![entity_id], |row| {
Ok(GraphEdge {
id: row.get(0)?,
from_id: row.get(1)?,
to_id: row.get(2)?,
edge_type: row.get(3)?,
data: serde_json::from_str(row.get_ref(4)?.as_str()?)
.map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?,
})
})?;
let mut edges = Vec::new();
for row in rows {
edges.push(row?);
}
Ok(edges)
})
}
fn get_outgoing_edges(graph: &SqliteGraph, entity_id: i64) -> Result<Vec<GraphEdge>> {
with_graph_conn(graph, |conn| {
let mut stmt = conn.prepare_cached(
"SELECT id, from_id, to_id, edge_type, data FROM graph_edges WHERE from_id=?1 ORDER BY id",
)?;
let rows = stmt.query_map(params![entity_id], |row| {
Ok(GraphEdge {
id: row.get(0)?,
from_id: row.get(1)?,
to_id: row.get(2)?,
edge_type: row.get(3)?,
data: serde_json::from_str(row.get_ref(4)?.as_str()?)
.map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?,
})
})?;
let mut edges = Vec::new();
for row in rows {
edges.push(row?);
}
Ok(edges)
})
}