Skip to main content

kaizen/extensions/
diffs.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2
3use crate::store::Store;
4use crate::store::tool_span_index::ToolSpanRecord;
5use anyhow::Result;
6use rusqlite::params;
7use serde::{Deserialize, Serialize};
8
9#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
10pub struct StepDiff {
11    pub session_id: String,
12    pub span_id: String,
13    pub files: Vec<String>,
14    pub added_lines: u64,
15    pub removed_lines: u64,
16    pub raw_patch_stored: bool,
17}
18
19pub fn upsert_tool_span(store: &Store, span: &ToolSpanRecord) -> Result<()> {
20    if span.paths.is_empty() {
21        return Ok(());
22    }
23    insert(
24        store,
25        &StepDiff {
26            session_id: span.session_id.clone(),
27            span_id: span.span_id.clone(),
28            files: span.paths.clone(),
29            added_lines: 0,
30            removed_lines: 0,
31            raw_patch_stored: false,
32        },
33    )
34}
35
36pub fn refresh_session(
37    store: &Store,
38    session_id: &str,
39    capture_raw_patch: bool,
40) -> Result<Vec<StepDiff>> {
41    let rows: Vec<StepDiff> = store
42        .tool_spans_for_session(session_id)?
43        .into_iter()
44        .filter(|span| !span.paths.is_empty())
45        .map(|span| StepDiff {
46            session_id: session_id.to_string(),
47            span_id: span.span_id,
48            files: span.paths,
49            added_lines: 0,
50            removed_lines: 0,
51            raw_patch_stored: capture_raw_patch,
52        })
53        .collect();
54    store
55        .conn()
56        .execute("DELETE FROM step_diffs WHERE session_id = ?1", [session_id])?;
57    rows.iter().try_for_each(|row| insert(store, row))?;
58    Ok(rows)
59}
60
61fn insert(store: &Store, row: &StepDiff) -> Result<()> {
62    let files_json = serde_json::to_string(&row.files)?;
63    store.conn().execute(
64        "INSERT OR REPLACE INTO step_diffs(
65            session_id, span_id, files_json, added_lines, removed_lines, raw_patch
66         ) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
67        params![
68            row.session_id,
69            row.span_id,
70            files_json,
71            row.added_lines as i64,
72            row.removed_lines as i64,
73            Option::<String>::None,
74        ],
75    )?;
76    Ok(())
77}