kaynine-core 0.1.0

use crate::error::RunFailureReason;
use crate::ids::ToolCallId;
use crate::message::{ContentBlock, Message, ToolResultPayload};
use crate::provider::ModelMessage;

#[derive(Clone, Debug, PartialEq, thiserror::Error)]
pub enum ChainError {
    #[error("dangling tool call {call_id}")]
    DanglingToolCall { call_id: ToolCallId },
    #[error("unexpected tool result {call_id}")]
    UnexpectedToolResult { call_id: ToolCallId },
    #[error("tool result order mismatch at position {position}")]
    ToolResultOrderMismatch { position: usize },
}

impl From<ChainError> for RunFailureReason {
    fn from(_: ChainError) -> Self {
        RunFailureReason::Internal
    }
}

fn unsatisfied(pending: &[ToolCallId], accumulated: usize) -> ChainError {
    ChainError::DanglingToolCall {
        call_id: pending
            .get(accumulated)
            .or_else(|| pending.first())
            .cloned()
            .expect("pending must be non-empty"),
    }
}

pub fn project_and_validate(history: &[Message]) -> Result<Vec<ModelMessage>, ChainError> {
    let mut projected: Vec<ModelMessage> = Vec::new();
    let mut pending: Vec<ToolCallId> = Vec::new();
    let mut accumulated: Vec<ToolResultPayload> = Vec::new();

    for message in history {
        match message {
            Message::User { blocks } => {
                if !pending.is_empty() {
                    return Err(unsatisfied(&pending, accumulated.len()));
                }
                projected.push(ModelMessage::User {
                    blocks: blocks.clone(),
                });
            }
            Message::Assistant { blocks, .. } => {
                if !pending.is_empty() {
                    return Err(unsatisfied(&pending, accumulated.len()));
                }
                projected.push(ModelMessage::Assistant {
                    blocks: blocks.clone(),
                });
                accumulated.clear();
                pending = blocks
                    .iter()
                    .filter_map(|b| match b {
                        ContentBlock::ToolCall { id, .. } => Some(id.clone()),
                        _ => None,
                    })
                    .collect();
            }
            Message::ToolResult { results } => {
                for result in results {
                    let position = accumulated.len();
                    if position >= pending.len() {
                        return Err(ChainError::UnexpectedToolResult {
                            call_id: result.call_id.clone(),
                        });
                    }
                    if result.call_id != pending[position] {
                        if !pending[position..].contains(&result.call_id) {
                            return Err(ChainError::UnexpectedToolResult {
                                call_id: result.call_id.clone(),
                            });
                        }
                        return Err(ChainError::ToolResultOrderMismatch { position });
                    }
                    accumulated.push(result.clone());
                }
                if accumulated.len() == pending.len() {
                    projected.push(ModelMessage::ToolResults {
                        results: std::mem::take(&mut accumulated),
                    });
                    pending.clear();
                }
            }
            Message::Summary { text } => {
                // A checkpoint replaces everything before it, so nothing can
                // be pending across it (validation would have failed first).
                if !pending.is_empty() {
                    return Err(unsatisfied(&pending, accumulated.len()));
                }
                projected.push(ModelMessage::Summary { text: text.clone() });
            }
        }
    }

    if !pending.is_empty() {
        return Err(unsatisfied(&pending, accumulated.len()));
    }
    Ok(projected)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ids::ToolCallId;
    use crate::message::{ContentBlock, FinishReason, Message, ToolResultPayload};
    use crate::provider::ModelMessage;

    fn user(text: &str) -> Message {
        Message::User {
            blocks: vec![ContentBlock::Text { text: text.into() }],
        }
    }

    fn assistant_with_calls(ids: &[&str]) -> Message {
        Message::Assistant {
            blocks: ids
                .iter()
                .map(|id| ContentBlock::ToolCall {
                    id: ToolCallId::from(*id),
                    name: "t".into(),
                    arguments: serde_json::json!({}),
                })
                .collect(),
            finish_reason: FinishReason::Stop,
            truncated: false,
        }
    }

    fn results(ids: &[&str]) -> Message {
        Message::ToolResult {
            results: ids
                .iter()
                .map(|id| ToolResultPayload {
                    call_id: ToolCallId::from(*id),
                    is_error: false,
                    text: "ok".into(),
                })
                .collect(),
        }
    }

    #[test]
    fn projects_valid_chain_in_order() {
        let history = vec![
            user("q"),
            assistant_with_calls(&["c1", "c2"]),
            results(&["c1", "c2"]),
        ];
        let projected = project_and_validate(&history).unwrap();
        assert_eq!(projected.len(), 3);
        assert!(matches!(projected[0], ModelMessage::User { .. }));
        assert!(matches!(projected[2], ModelMessage::ToolResults { .. }));
    }

    #[test]
    fn rejects_dangling_tool_call_at_end() {
        let history = vec![user("q"), assistant_with_calls(&["c1"])];
        assert!(matches!(
            project_and_validate(&history),
            Err(ChainError::DanglingToolCall { call_id }) if call_id == ToolCallId::from("c1")
        ));
    }

    #[test]
    fn rejects_dangling_tool_call_before_next_assistant() {
        let history = vec![
            user("q"),
            assistant_with_calls(&["c1"]),
            results(&["c1"]),
            assistant_with_calls(&["c2"]),
            results(&["c2"]),
            assistant_with_calls(&["c3"]),
            results(&["c3"]),
            assistant_with_calls(&["c4"]),
        ];
        assert!(matches!(
            project_and_validate(&history),
            Err(ChainError::DanglingToolCall { call_id }) if call_id == ToolCallId::from("c4")
        ));
    }

    #[test]
    fn rejects_unexpected_tool_result() {
        let history = vec![
            user("q"),
            assistant_with_calls(&["c1"]),
            results(&["c1", "c2"]),
        ];
        assert!(matches!(
            project_and_validate(&history),
            Err(ChainError::UnexpectedToolResult { call_id }) if call_id == ToolCallId::from("c2")
        ));
    }

    #[test]
    fn rejects_out_of_order_results() {
        let history = vec![
            user("q"),
            assistant_with_calls(&["c1", "c2"]),
            results(&["c2", "c1"]),
        ];
        assert!(matches!(
            project_and_validate(&history),
            Err(ChainError::ToolResultOrderMismatch { position: 0, .. })
        ));
    }

    #[test]
    fn rejects_missing_results_in_batch() {
        let history = vec![
            user("q"),
            assistant_with_calls(&["c1", "c2"]),
            results(&["c1"]),
        ];
        assert!(matches!(
            project_and_validate(&history),
            Err(ChainError::DanglingToolCall { call_id }) if call_id == ToolCallId::from("c2")
        ));
    }

    #[test]
    fn rejects_tool_result_without_pending_calls() {
        let history = vec![user("q"), results(&["c9"])];
        assert!(matches!(
            project_and_validate(&history),
            Err(ChainError::UnexpectedToolResult { call_id }) if call_id == ToolCallId::from("c9")
        ));
    }

    #[test]
    fn consecutive_per_call_results_group_into_one_batch() {
        let history = vec![
            user("q"),
            assistant_with_calls(&["c1", "c2"]),
            results(&["c1"]),
            results(&["c2"]),
        ];
        let projected = project_and_validate(&history).unwrap();
        assert_eq!(projected.len(), 3);
        assert!(matches!(projected[0], ModelMessage::User { .. }));
        assert!(matches!(projected[1], ModelMessage::Assistant { .. }));
        match &projected[2] {
            ModelMessage::ToolResults { results } => {
                assert_eq!(results.len(), 2);
                assert_eq!(results[0].call_id, ToolCallId::from("c1"));
                assert_eq!(results[1].call_id, ToolCallId::from("c2"));
            }
            other => panic!("expected ToolResults, got {other:?}"),
        }
    }

    #[test]
    fn interleaved_out_of_order_per_call_results_rejected() {
        let history = vec![
            user("q"),
            assistant_with_calls(&["c1", "c2"]),
            results(&["c2"]),
            results(&["c1"]),
        ];
        assert!(matches!(
            project_and_validate(&history),
            Err(ChainError::ToolResultOrderMismatch { position: 0 })
        ));
    }

    #[test]
    fn incomplete_group_before_next_assistant_is_dangling() {
        let history = vec![
            user("q"),
            assistant_with_calls(&["c1", "c2"]),
            results(&["c1"]),
            assistant_with_calls(&["c3"]),
        ];
        assert!(matches!(
            project_and_validate(&history),
            Err(ChainError::DanglingToolCall { call_id }) if call_id == ToolCallId::from("c2")
        ));
    }

    #[test]
    fn per_call_results_for_three_calls() {
        let history = vec![
            user("q"),
            assistant_with_calls(&["c1", "c2", "c3"]),
            results(&["c1"]),
            results(&["c2"]),
            results(&["c3"]),
        ];
        let projected = project_and_validate(&history).unwrap();
        assert_eq!(projected.len(), 3);
        match &projected[2] {
            ModelMessage::ToolResults { results } => assert_eq!(results.len(), 3),
            other => panic!("expected ToolResults, got {other:?}"),
        }
    }

    #[test]
    fn empty_tool_result_batch_with_pending_is_dangling() {
        let history = vec![
            user("q"),
            assistant_with_calls(&["c1"]),
            Message::ToolResult {
                results: Vec::new(),
            },
        ];
        assert!(matches!(
            project_and_validate(&history),
            Err(ChainError::DanglingToolCall { call_id }) if call_id == ToolCallId::from("c1")
        ));
    }
}