agent-runtime-kernel 0.1.1

use std::collections::BTreeMap;

use runtime_ports::{ModelContent, ModelMessage};
use sha2::{Digest, Sha256};

const STUB_PREFIX: &str = "Runtime repeated failure elided: signature=";
const COUNT_PREFIX: &str = "Runtime failure repeat count:";

#[derive(Default)]
struct FailureGroup {
    full: Vec<(usize, usize)>,
    stubs: Vec<(usize, usize)>,
}

/// Fold repeated failed observations without removing their Tool messages, so
/// every assistant ToolUse keeps its matching ToolResult transaction.
pub(crate) fn fold_repeated_failures(messages: &mut [ModelMessage]) {
    let mut groups = BTreeMap::<String, FailureGroup>::new();
    for (message_index, message) in messages.iter().enumerate() {
        for (content_index, content) in message.content.iter().enumerate() {
            let ModelContent::ToolResult {
                name,
                content,
                failed: true,
                ..
            } = content
            else {
                continue;
            };
            if let Some(signature) = stub_signature(content) {
                groups
                    .entry(signature.to_string())
                    .or_default()
                    .stubs
                    .push((message_index, content_index));
                continue;
            }
            let Some(signature) = failure_signature(name, content) else {
                continue;
            };
            groups
                .entry(signature)
                .or_default()
                .full
                .push((message_index, content_index));
        }
    }

    for (signature, group) in groups {
        let total = group.full.len().saturating_add(group.stubs.len());
        if total <= 1 {
            if let Some(location) = group.full.first().copied() {
                rewrite(messages, location, strip_count_header);
            }
            continue;
        }
        let latest = group.full.iter().max().copied();
        for location in group
            .full
            .iter()
            .copied()
            .filter(|location| Some(*location) != latest)
            .chain(group.stubs.iter().copied())
        {
            let stub = format!(
                "{STUB_PREFIX}{signature}; occurrences={total}. The latest full instance is kept later in context."
            );
            rewrite(messages, location, |_| stub);
        }
        if let Some(location) = latest {
            rewrite(messages, location, |content| {
                let content = strip_count_header(content);
                format!(
                    "{COUNT_PREFIX} this same failure has occurred {total} times; older duplicates are summarized.\n{content}"
                )
            });
        }
    }
}

fn rewrite(
    messages: &mut [ModelMessage],
    (message_index, content_index): (usize, usize),
    update: impl FnOnce(&str) -> String,
) {
    let Some(ModelContent::ToolResult { content, .. }) = messages
        .get_mut(message_index)
        .and_then(|message| message.content.get_mut(content_index))
    else {
        return;
    };
    *content = update(content);
}

fn stub_signature(content: &str) -> Option<&str> {
    content.strip_prefix(STUB_PREFIX)?.split(';').next()
}

fn strip_count_header(content: &str) -> String {
    content
        .lines()
        .filter(|line| !line.starts_with(COUNT_PREFIX))
        .collect::<Vec<_>>()
        .join("\n")
}

fn failure_signature(name: &str, content: &str) -> Option<String> {
    let content = strip_count_header(content);
    let has_marker = content.lines().any(|line| {
        let lower = line.to_ascii_lowercase();
        lower.contains("error:")
            || lower.contains("failed")
            || lower.contains("failure")
            || lower.contains("panicked at")
            || nonzero_exit(line)
    });
    if !has_marker {
        return None;
    }
    let mut normalized = name.trim().to_string();
    for line in content
        .lines()
        .map(str::trim)
        .filter(|line| !line.is_empty())
        .take(40)
    {
        normalized.push('\n');
        normalized.push_str(&line.split_whitespace().collect::<Vec<_>>().join(" "));
    }
    let digest = Sha256::digest(normalized.as_bytes());
    Some(format!("{:x}", digest))
}

fn nonzero_exit(line: &str) -> bool {
    line.trim()
        .strip_prefix("Exit code:")
        .and_then(|value| value.trim().parse::<i32>().ok())
        .is_some_and(|code| code != 0)
}

#[cfg(test)]
mod tests {
    use runtime_ports::{ModelContent, ModelMessage, ModelRole};

    use super::*;

    fn failed(id: &str, text: &str) -> ModelMessage {
        ModelMessage {
            role: ModelRole::Tool,
            content: vec![ModelContent::ToolResult {
                tool_use_id: id.into(),
                name: "run".into(),
                content: text.into(),
                failed: true,
            }],
        }
    }

    #[test]
    fn keeps_latest_full_failure_and_stubs_older_occurrences_idempotently() {
        let mut messages = vec![
            failed("one", "Command: test\nExit code: 1\nError: boom"),
            failed("two", "Command: test\nExit code: 1\nError: boom"),
        ];
        fold_repeated_failures(&mut messages);
        fold_repeated_failures(&mut messages);
        let encoded = serde_json::to_string(&messages).unwrap();
        assert!(encoded.contains(STUB_PREFIX));
        assert!(encoded.contains("occurred 2 times"));
        assert_eq!(messages.len(), 2);
    }
}