use crate::context::ContextReader;
use crate::message::{
AssistantResponse, Context, ContextListing, ConversationMessage, MessageKind, TextContent,
ToolCall, ToolResultData,
};
use anyhow::Context as _;
use rusqlite::Connection;
use serde_json::Value;
use std::path::PathBuf;
pub struct CrushReader {
db_path: PathBuf,
}
impl CrushReader {
pub fn new(db_path: PathBuf) -> Self {
Self { db_path }
}
}
impl ContextReader for CrushReader {
fn list_contexts(&self) -> anyhow::Result<Vec<ContextListing>> {
let conn = Connection::open(&self.db_path)
.with_context(|| format!("failed to open crush db at {}", self.db_path.display()))?;
let sql = "SELECT session_id, COUNT(*) FROM messages GROUP BY session_id";
let mut stmt = conn.prepare(sql)?;
let rows = stmt.query_map([], |row| {
Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)?))
})?;
let mut contexts = Vec::new();
for row in rows {
let (session_id, count) = row?;
let detail = format!("{count} messages");
contexts.push(ContextListing {
id: session_id,
detail,
path: Some(self.db_path.clone()),
});
}
Ok(contexts)
}
fn read_context(&self, context_id: &str) -> anyhow::Result<Context> {
let conn = Connection::open(&self.db_path)
.with_context(|| format!("failed to open crush db at {}", self.db_path.display()))?;
let sql = "SELECT role, parts FROM messages WHERE session_id = ?1 ORDER BY created_at";
let mut stmt = conn.prepare(sql)?;
let rows = stmt.query_map(rusqlite::params![context_id], |row| {
Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
})?;
let mut messages = Vec::new();
for (idx, row) in rows.enumerate() {
let (role, parts_str) = row?;
let parts: Value = serde_json::from_str(&parts_str)?;
let crush_parts = crush_build_message_kind(&role, &parts)?;
let entry_id = format!("crush-{idx}");
messages.push(ConversationMessage {
entry_id: entry_id.clone(),
kind: crush_parts.kind,
});
for (tri, tr) in crush_parts.tool_results.into_iter().enumerate() {
let tr_entry_id = format!("crush-{idx}-tr{tri}");
messages.push(ConversationMessage {
entry_id: tr_entry_id,
kind: MessageKind::ToolResultData(tr),
});
}
}
Ok(Context {
entries: Vec::new(),
messages,
})
}
}
struct CrushMessageParts {
kind: MessageKind,
tool_results: Vec<ToolResultData>,
}
fn crush_build_message_kind(role: &str, parts: &Value) -> anyhow::Result<CrushMessageParts> {
let parts = parts.as_array().context("parts is not an array")?;
match role {
"user" => {
let text: String = parts
.iter()
.filter(|p| p["type"].as_str() == Some("text"))
.map(|p| p["text"].as_str().unwrap_or("").to_string())
.collect::<Vec<_>>()
.join("\n");
Ok(CrushMessageParts {
kind: MessageKind::TextContent(TextContent {
role: "user".to_string(),
text,
}),
tool_results: Vec::new(),
})
}
"assistant" => {
let mut text = String::new();
let mut tool_calls = Vec::new();
let mut tool_results = Vec::new();
for part in parts {
match part["type"].as_str() {
Some("text") => {
if let Some(t) = part["text"].as_str() {
text.push_str(t);
}
}
Some("tool_use") => {
let name = part["name"].as_str().unwrap_or("").to_string();
let arguments = part.get("input").cloned().unwrap_or(Value::Null);
tool_calls.push(ToolCall {
name: name.clone(),
arguments,
});
if let Some(res) = part["result"].as_str() {
tool_results.push(ToolResultData {
tool_name: name,
content: res.to_string(),
is_error: false,
});
}
}
Some("tool_result") => {
let content = part["content"].as_str().unwrap_or("").to_string();
let tool_name = part["name"].as_str().map_or_else(
|| "unknown".to_string(),
std::string::ToString::to_string,
);
let is_error = part["is_error"].as_bool().unwrap_or(false);
tool_results.push(ToolResultData {
tool_name,
content,
is_error,
});
}
_ => {}
}
}
Ok(CrushMessageParts {
kind: MessageKind::AssistantResponse(AssistantResponse {
thinking: Vec::new(),
tool_calls,
text,
}),
tool_results,
})
}
_ => {
let text: String = parts
.iter()
.filter(|p| p["type"].as_str() == Some("text"))
.map(|p| p["text"].as_str().unwrap_or("").to_string())
.collect::<Vec<_>>()
.join("\n");
Ok(CrushMessageParts {
kind: MessageKind::TextContent(TextContent {
role: role.to_string(),
text,
}),
tool_results: Vec::new(),
})
}
}
}