haematite-cli 0.21.0

use std::num::NonZeroU32;

use serde::{Deserialize, Serialize};
use serde_json::value::RawValue;

use super::read_support::{Snapshot, bounded, call, failure};
use super::{EncodeError, Sender, ServiceRefusal, ToolResult, refused};
use crate::mcp::{rpc, types};
use crate::serve::protocol;

#[derive(Deserialize)]
struct Capture {
    selection: Snapshot,
}

#[derive(Deserialize)]
struct Document {
    id: String,
    raw_json: String,
}

#[derive(Deserialize)]
struct Record {
    selection: Snapshot,
    document: Option<Document>,
}

#[derive(Deserialize)]
struct History {
    entries: Vec<types::HistoryEntry>,
    next_after: Option<types::Version>,
}

#[derive(Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
struct Cursor {
    selection: types::Snapshot,
    after: types::Version,
}

fn after(input: &types::ReadHistoryInput) -> Result<Option<types::Version>, ServiceRefusal> {
    let Some(raw) = input.after.as_deref() else {
        return Ok(None);
    };
    let version = input.selection.version.ok_or_else(|| {
        failure(
            "mcp_version_required",
            "History continuation requires the captured selection version",
        )
    })?;
    let cursor: Cursor = serde_json::from_str(raw)
        .map_err(|error| failure("mcp_cursor_invalid", error.to_string()))?;
    if cursor.selection.branch != input.selection.branch
        || cursor.selection.version.0 != version.0
        || cursor.after.0.generation != version.0.generation
        || cursor.after.0.sequence > version.0.sequence
    {
        return Err(failure(
            "mcp_cursor_invalid",
            "History continuation is beyond the captured generation or version",
        ));
    }
    Ok(Some(cursor.after))
}

fn capture(
    sender: &mut impl Sender,
    input: &types::ReadHistoryInput,
) -> Result<(Snapshot, Option<types::RecordValue>), ServiceRefusal> {
    let selected = protocol::Selection {
        branch: input.selection.branch.clone(),
        version: input.selection.version.map(|version| version.0),
    };
    if let Some(record) = &input.record {
        let answer: Record = call(
            sender,
            protocol::Operation::Get(protocol::DocumentRead {
                selection: selected,
                collection: record.collection.clone(),
                document_id: record.document_id.clone(),
            }),
        )?;
        answer.selection.requested(&input.selection)?;
        let value = answer
            .document
            .map(|body| {
                if body.id != record.document_id {
                    return Err(failure(
                        "mcp_service_answer_invalid",
                        "The served document does not identify the requested record",
                    ));
                }
                let document = RawValue::from_string(body.raw_json)
                    .map_err(|error| failure("mcp_service_answer_invalid", error.to_string()))?;
                Ok(types::RecordValue {
                    record: types::RecordRef {
                        collection: record.collection.clone(),
                        document_id: body.id,
                    },
                    document,
                })
            })
            .transpose()?;
        Ok((answer.selection, value))
    } else {
        let answer: Capture = call(
            sender,
            protocol::Operation::Collections(protocol::CollectionPage {
                selection: selected,
                after: None,
                limit: NonZeroU32::MIN,
            }),
        )?;
        answer.selection.requested(&input.selection)?;
        Ok((answer.selection, None))
    }
}

fn history(
    sender: &mut impl Sender,
    input: &types::ReadHistoryInput,
) -> Result<ToolResult, ServiceRefusal> {
    let limit = input.page.rows;
    let after = after(input)?;
    let (selected, record) = capture(sender, input)?;
    let history: History = call(
        sender,
        protocol::Operation::History(protocol::HistoryPage {
            branch: selected.branch.clone(),
            after: after.map(|version| version.0),
            limit,
        }),
    )?;
    let mut entries = Vec::new();
    for entry in history.entries {
        if entry.branch != selected.branch
            || entry.version.0.generation != selected.version.0.generation
        {
            return Err(failure(
                "mcp_service_snapshot_mismatch",
                "History does not identify the captured branch generation",
            ));
        }
        if entry.version.0.sequence <= selected.version.0.sequence {
            entries.push(entry);
        }
    }
    if history
        .next_after
        .is_some_and(|after| after.0.generation != selected.version.0.generation)
    {
        return Err(failure(
            "mcp_service_answer_invalid",
            "The history cursor identifies another branch generation",
        ));
    }
    let next = history
        .next_after
        .filter(|after| after.0.sequence < selected.version.0.sequence);
    let next_after = next
        .map(|after| {
            serde_json::to_string(&Cursor {
                selection: types::Snapshot {
                    branch: selected.branch.clone(),
                    version: selected.version,
                },
                after,
            })
            .map_err(|error| failure("mcp_cursor_encoding_failed", error.to_string()))
        })
        .transpose()?;
    let output = types::ReadHistoryOutput {
        selection: selected.output(),
        history: types::Page {
            entries,
            next_after,
        },
        record,
    };
    bounded(&output, input.page.bytes.get())
}

pub(super) fn dispatch(
    sender: &mut impl Sender,
    request: &rpc::ToolCall,
) -> Result<ToolResult, EncodeError> {
    let input = match serde_json::from_str(request.arguments.get()) {
        Ok(input) => input,
        Err(error) => return refused(&failure("mcp_arguments_invalid", error.to_string())),
    };
    match history(sender, &input) {
        Ok(answer) => Ok(answer),
        Err(error) => refused(&error),
    }
}