deepstrike_core/memory/

trace_analyzer.rs

use std::collections::HashMap;

use crate::types::message::{Content, ContentPart, Message, Role};

/// A single insight extracted from a session's message trace.
#[derive(Debug, Clone)]
pub struct TraceInsight {
    pub kind: InsightKind,
    /// 0.0–1.0; scales with evidence strength (error frequency, sequence length, etc.)
    pub confidence: f64,
    pub session_id: String,
}

#[derive(Debug, Clone)]
pub enum InsightKind {
    /// The same tool failed at least `min_error_count` times in one session.
    RepeatedToolError {
        tool_name: String,
        error_count: usize,
        /// First observed error output, truncated to 200 chars.
        sample_error: String,
    },
    /// A consecutive run of tool calls that completed without any error,
    /// followed by an assistant turn with no further tool calls.
    SuccessfulToolSequence {
        tools: Vec<String>,
        /// The user message that kicked off the sequence, truncated to 100 chars.
        context_hint: String,
    },
    /// An assistant message whose reasoning block exceeds `min_reasoning_chars`.
    LongReasoning {
        /// First 300 chars of the assistant text.
        summary_hint: String,
    },
    /// Free-form insight generated by the LLM during the dreaming synthesis phase.
    Synthesized { text: String },
}

impl InsightKind {
    pub fn tag(&self) -> &'static str {
        match self {
            Self::RepeatedToolError { .. } => "repeated_tool_error",
            Self::SuccessfulToolSequence { .. } => "successful_sequence",
            Self::LongReasoning { .. } => "long_reasoning",
            Self::Synthesized { .. } => "synthesized",
        }
    }
}

#[derive(Debug, Clone)]
pub struct AnalysisPolicy {
    /// Minimum number of errors on the same tool before flagging. Default: 2.
    pub min_error_count: usize,
    /// Minimum consecutive tool calls for a sequence to be notable. Default: 2.
    pub min_success_sequence_len: usize,
    /// Minimum char length of an assistant message to treat as long reasoning. Default: 500.
    pub min_reasoning_chars: usize,
}

impl Default for AnalysisPolicy {
    fn default() -> Self {
        Self {
            min_error_count: 2,
            min_success_sequence_len: 2,
            min_reasoning_chars: 500,
        }
    }
}

pub struct TraceAnalyzer {
    pub policy: AnalysisPolicy,
}

impl TraceAnalyzer {
    pub fn new(policy: AnalysisPolicy) -> Self {
        Self { policy }
    }

    /// Analyze all sessions in one pass. Sessions are `(session_id, messages)` tuples.
    pub fn analyze_batch(&self, sessions: &[(String, Vec<Message>)]) -> Vec<TraceInsight> {
        sessions
            .iter()
            .flat_map(|(id, msgs)| self.analyze(id, msgs))
            .collect()
    }

    pub fn analyze(&self, session_id: &str, messages: &[Message]) -> Vec<TraceInsight> {
        let mut insights = Vec::new();
        insights.extend(self.detect_repeated_errors(session_id, messages));
        insights.extend(self.detect_successful_sequences(session_id, messages));
        insights.extend(self.detect_long_reasoning(session_id, messages));
        insights
    }

    // --- private detectors ---------------------------------------------------

    fn detect_repeated_errors(&self, session_id: &str, messages: &[Message]) -> Vec<TraceInsight> {
        // Build call_id → tool_name index from assistant messages.
        let mut call_id_to_name: HashMap<String, String> = HashMap::new();
        for msg in messages {
            if msg.role == Role::Assistant {
                for tc in &msg.tool_calls {
                    call_id_to_name.insert(tc.id.to_string(), tc.name.to_string());
                }
            }
        }

        // Count errors per tool; store first observed sample.
        let mut error_counts: HashMap<String, (usize, String)> = HashMap::new();
        for msg in messages {
            if msg.role != Role::Tool {
                continue;
            }
            if let Content::Parts(parts) = &msg.content {
                for part in parts {
                    if let ContentPart::ToolResult {
                        call_id,
                        output,
                        is_error,
                    } = part
                    {
                        if *is_error {
                            if let Some(name) = call_id_to_name.get(call_id.as_str()) {
                                let entry = error_counts
                                    .entry(name.clone())
                                    .or_insert_with(|| (0, output.chars().take(200).collect()));
                                entry.0 += 1;
                            }
                        }
                    }
                }
            }
        }

        error_counts
            .into_iter()
            .filter(|(_, (count, _))| *count >= self.policy.min_error_count)
            .map(|(tool_name, (error_count, sample_error))| TraceInsight {
                kind: InsightKind::RepeatedToolError {
                    tool_name,
                    error_count,
                    sample_error,
                },
                // Confidence scales with frequency, saturating at 1.0 after 5 errors.
                confidence: (error_count as f64 / 5.0).min(1.0),
                session_id: session_id.to_string(),
            })
            .collect()
    }

    fn detect_successful_sequences(
        &self,
        session_id: &str,
        messages: &[Message],
    ) -> Vec<TraceInsight> {
        let mut insights = Vec::new();
        let mut sequence: Vec<String> = Vec::new();
        let mut context_hint = String::new();
        let mut sequence_has_error = false;

        for msg in messages {
            match msg.role {
                Role::User => {
                    // The most recent user message becomes the context hint for the next sequence.
                    if let Some(text) = msg.content.as_text() {
                        context_hint = text.chars().take(100).collect();
                    }
                }
                Role::Assistant => {
                    if msg.tool_calls.is_empty() {
                        // Assistant is done — emit insight if the sequence was clean.
                        if !sequence_has_error
                            && sequence.len() >= self.policy.min_success_sequence_len
                        {
                            let confidence =
                                (sequence.len() as f64 / 10.0).min(0.9_f64).max(0.5_f64);
                            insights.push(TraceInsight {
                                kind: InsightKind::SuccessfulToolSequence {
                                    tools: sequence.clone(),
                                    context_hint: context_hint.clone(),
                                },
                                confidence,
                                session_id: session_id.to_string(),
                            });
                        }
                        sequence.clear();
                        sequence_has_error = false;
                    } else {
                        for tc in &msg.tool_calls {
                            sequence.push(tc.name.to_string());
                        }
                    }
                }
                Role::Tool => {
                    if let Content::Parts(parts) = &msg.content {
                        if parts
                            .iter()
                            .any(|p| matches!(p, ContentPart::ToolResult { is_error: true, .. }))
                        {
                            sequence_has_error = true;
                        }
                    }
                }
                Role::System => {}
            }
        }

        insights
    }

    fn detect_long_reasoning(&self, session_id: &str, messages: &[Message]) -> Vec<TraceInsight> {
        messages
            .iter()
            .filter(|m| m.role == Role::Assistant)
            .filter_map(|m| m.content.as_text())
            .filter(|text| text.len() >= self.policy.min_reasoning_chars)
            .map(|text| {
                let summary_hint: String = text.chars().take(300).collect();
                // Confidence grows with length, saturating at 0.8 around 2000 chars.
                let confidence = (text.len() as f64 / 2000.0).min(0.8_f64).max(0.4_f64);
                TraceInsight {
                    kind: InsightKind::LongReasoning { summary_hint },
                    confidence,
                    session_id: session_id.to_string(),
                }
            })
            .collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::message::{ContentPart, ToolCall};
    use compact_str::CompactString;
    use pretty_assertions::assert_eq;

    fn analyzer() -> TraceAnalyzer {
        TraceAnalyzer::new(AnalysisPolicy::default())
    }

    fn assistant_with_tool(call_id: &str, tool_name: &str) -> Message {
        let mut msg = Message::assistant("");
        msg.tool_calls = vec![ToolCall {
            id: CompactString::new(call_id),
            name: CompactString::new(tool_name),
            arguments: serde_json::Value::Null,
        }];
        msg
    }

    fn tool_error(call_id: &str, err: &str) -> Message {
        Message::tool(vec![ContentPart::ToolResult {
            call_id: CompactString::new(call_id),
            output: err.to_string(),
            is_error: true,
        }])
    }

    fn tool_ok(call_id: &str) -> Message {
        Message::tool(vec![ContentPart::ToolResult {
            call_id: CompactString::new(call_id),
            output: "ok".to_string(),
            is_error: false,
        }])
    }

    #[test]
    fn detects_repeated_tool_errors() {
        let messages = vec![
            assistant_with_tool("c1", "bash"),
            tool_error("c1", "permission denied"),
            assistant_with_tool("c2", "bash"),
            tool_error("c2", "permission denied"),
        ];
        let insights = analyzer().analyze("s1", &messages);
        let errors: Vec<_> = insights
            .iter()
            .filter(|i| matches!(i.kind, InsightKind::RepeatedToolError { .. }))
            .collect();
        assert_eq!(errors.len(), 1);
        if let InsightKind::RepeatedToolError {
            tool_name,
            error_count,
            ..
        } = &errors[0].kind
        {
            assert_eq!(tool_name, "bash");
            assert_eq!(*error_count, 2);
        }
    }

    #[test]
    fn skips_single_error_below_threshold() {
        let messages = vec![assistant_with_tool("c1", "bash"), tool_error("c1", "oops")];
        let insights = analyzer().analyze("s1", &messages);
        assert!(
            insights
                .iter()
                .all(|i| !matches!(i.kind, InsightKind::RepeatedToolError { .. }))
        );
    }

    #[test]
    fn detects_successful_tool_sequence() {
        let messages = vec![
            Message::user("fix the bug"),
            assistant_with_tool("c1", "read_file"),
            tool_ok("c1"),
            assistant_with_tool("c2", "edit_file"),
            tool_ok("c2"),
            Message::assistant("Done!"),
        ];
        let insights = analyzer().analyze("s1", &messages);
        let seqs: Vec<_> = insights
            .iter()
            .filter(|i| matches!(i.kind, InsightKind::SuccessfulToolSequence { .. }))
            .collect();
        assert_eq!(seqs.len(), 1);
        if let InsightKind::SuccessfulToolSequence {
            tools,
            context_hint,
        } = &seqs[0].kind
        {
            assert_eq!(tools, &["read_file", "edit_file"]);
            assert!(context_hint.contains("fix the bug"));
        }
    }

    #[test]
    fn resets_sequence_on_error() {
        let messages = vec![
            Message::user("do something"),
            assistant_with_tool("c1", "bash"),
            tool_error("c1", "fail"),
            assistant_with_tool("c2", "bash"),
            tool_ok("c2"),
            Message::assistant("Done"),
        ];
        let insights = analyzer().analyze("s1", &messages);
        // Sequence only has 1 clean tool call after the error, below min_success_sequence_len=2.
        assert!(
            insights
                .iter()
                .all(|i| !matches!(i.kind, InsightKind::SuccessfulToolSequence { .. }))
        );
    }

    #[test]
    fn detects_long_reasoning() {
        let long_text = "a".repeat(600);
        let messages = vec![Message::assistant(long_text)];
        let insights = analyzer().analyze("s1", &messages);
        assert!(
            insights
                .iter()
                .any(|i| matches!(i.kind, InsightKind::LongReasoning { .. }))
        );
    }
}