haematite-cli 0.21.7

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

use super::{EncodeError, Sender, ServiceRefusal, ToolResult, refused, result};
use crate::mcp::{rpc, types};
use crate::serve::protocol;

#[derive(Deserialize)]
struct Receipt {
    branch: String,
    version: types::Version,
    request_id: String,
}

fn arguments(error: &serde_json::Error) -> ServiceRefusal {
    ServiceRefusal {
        code: "mcp_arguments_invalid".to_owned(),
        message: error.to_string(),
        receipt: None,
    }
}

fn unknown_receipt(request_id: &str, problem: impl std::fmt::Display) -> ServiceRefusal {
    ServiceRefusal {
        code: "mcp_client_outcome_unknown".to_owned(),
        message: format!(
            "mcp_service_answer_invalid: {problem}; mutation.apply request_id {request_id} was sent; reconcile through mutation.receipt before further mutations"
        ),
        receipt: None,
    }
}

/// The collection the assistant recorder alone writes.
///
/// It holds the transcript and permission decisions a person reads back. A harness that reached it through an MCP tool could write its own approval, so every public tool call naming it
/// is refused before it reaches the store.
pub const ASSISTANT_RECORD: &str = "_assistant";

/// Who a write tool call comes from: the public MCP surfaces, or the assistant recorder in the door.
#[derive(Clone, Copy, PartialEq, Eq)]
pub enum Writer {
    Public,
    AssistantRecorder,
}

fn reserved(change: &protocol::Change) -> bool {
    let named = |collection: &str| collection == ASSISTANT_RECORD;
    match change {
        protocol::Change::Collection(write) => named(&write.collection),
        protocol::Change::PutMany(import) => named(&import.collection),
        protocol::Change::Delete(delete) => named(&delete.collection),
        protocol::Change::EdgePut(edge) => {
            named(&edge.from.collection) || named(&edge.to.collection)
        }
        protocol::Change::EdgeDelete(edge) => {
            named(&edge.from.collection) || named(&edge.to.collection)
        }
        _ => false,
    }
}

fn apply<Input: DeserializeOwned, Output: Serialize>(
    sender: &mut impl Sender,
    writer: Writer,
    call: &rpc::ToolCall,
    request: impl FnOnce(Input) -> Result<protocol::Mutation, ServiceRefusal>,
    output: impl FnOnce(Receipt) -> Output,
) -> Result<ToolResult, EncodeError> {
    let input: Input = match serde_json::from_str(call.arguments.get()) {
        Ok(input) => input,
        Err(error) => return refused(&arguments(&error)),
    };
    let request = match request(input) {
        Ok(request) => request,
        Err(refusal) => return refused(&refusal),
    };
    if writer == Writer::Public && reserved(&request.change) {
        return refused(&ServiceRefusal {
            code: "assistant_record_reserved".to_owned(),
            message: format!(
                "{ASSISTANT_RECORD} is written only by the assistant recorder; no MCP tool may write, delete, define or link it"
            ),
            receipt: None,
        });
    }
    let request_id = request.request_id.clone();
    let destination = match &request.change {
        protocol::Change::Fork(fork) => fork.name.clone(),
        _ => request.branch.clone(),
    };
    let raw = match sender.send(protocol::Operation::Mutate(request)) {
        Ok(raw) => raw,
        Err(refusal) => return refused(&refusal),
    };
    let receipt: Receipt = match serde_json::from_str(raw.get()) {
        Ok(receipt) => receipt,
        Err(error) => {
            return refused(&unknown_receipt(&request_id, error));
        }
    };
    if receipt.request_id != request_id || receipt.branch != destination {
        return refused(&unknown_receipt(
            &request_id,
            "The receipt does not identify the requested change and destination",
        ));
    }
    result(&output(receipt), false)
}

pub(super) fn define_collection(
    sender: &mut impl Sender,
    writer: Writer,
    call: &rpc::ToolCall,
) -> Result<ToolResult, EncodeError> {
    apply(
        sender,
        writer,
        call,
        |input: types::DefineCollectionInput| {
            let schema = match input.schema {
                types::CollectionShape::Remove => None,
                types::CollectionShape::Boolean(value) => Some(serde_json::Value::Bool(value)),
                types::CollectionShape::Rules(raw) => {
                    Some(serde_json::from_str(raw.get()).map_err(|error| arguments(&error))?)
                }
            };
            Ok(protocol::Mutation {
                request_id: input.request_id,
                branch: input.branch,
                expected: input.expected.0,
                change: protocol::Change::Collection(protocol::CollectionWrite {
                    collection: input.collection,
                    schema,
                }),
            })
        },
        |receipt| types::DefineCollectionOutput {
            branch: receipt.branch,
            version: receipt.version,
            request_id: receipt.request_id,
        },
    )
}

pub(super) fn put_records(
    sender: &mut impl Sender,
    writer: Writer,
    call: &rpc::ToolCall,
) -> Result<ToolResult, EncodeError> {
    apply(
        sender,
        writer,
        call,
        |input: types::PutRecordsInput| {
            Ok(protocol::Mutation {
                request_id: input.request_id,
                branch: input.branch,
                expected: input.expected.0,
                change: protocol::Change::PutMany(protocol::DocumentImport {
                    collection: input.collection,
                    documents: input
                        .records
                        .into_iter()
                        .map(|record| protocol::ImportedDocument {
                            document_id: record.document_id,
                            document: record.document,
                        })
                        .collect(),
                }),
            })
        },
        |receipt| types::PutRecordsOutput {
            branch: receipt.branch,
            version: receipt.version,
            request_id: receipt.request_id,
        },
    )
}

pub(super) fn delete_record(
    sender: &mut impl Sender,
    writer: Writer,
    call: &rpc::ToolCall,
) -> Result<ToolResult, EncodeError> {
    apply(
        sender,
        writer,
        call,
        |input: types::DeleteRecordInput| {
            Ok(protocol::Mutation {
                request_id: input.request_id,
                branch: input.branch,
                expected: input.expected.0,
                change: protocol::Change::Delete(protocol::DocumentDelete {
                    collection: input.record.collection,
                    document_id: input.record.document_id,
                }),
            })
        },
        |receipt| types::DeleteRecordOutput {
            branch: receipt.branch,
            version: receipt.version,
            request_id: receipt.request_id,
        },
    )
}

pub(super) fn link_records(
    sender: &mut impl Sender,
    writer: Writer,
    call: &rpc::ToolCall,
) -> Result<ToolResult, EncodeError> {
    apply(
        sender,
        writer,
        call,
        |input: types::LinkRecordsInput| {
            let properties: Box<RawValue> = serde_json::value::to_raw_value(&input.properties)
                .map_err(|error| arguments(&error))?;
            Ok(protocol::Mutation {
                request_id: input.request_id,
                branch: input.branch,
                expected: input.expected.0,
                change: protocol::Change::EdgePut(protocol::Edge {
                    relation: input.relation,
                    from: node(input.from),
                    to: node(input.to),
                    properties: serde_json::from_str(properties.get())
                        .map_err(|error| arguments(&error))?,
                }),
            })
        },
        |receipt| types::LinkRecordsOutput {
            branch: receipt.branch,
            version: receipt.version,
            request_id: receipt.request_id,
        },
    )
}

pub(super) fn unlink_records(
    sender: &mut impl Sender,
    writer: Writer,
    call: &rpc::ToolCall,
) -> Result<ToolResult, EncodeError> {
    apply(
        sender,
        writer,
        call,
        |input: types::UnlinkRecordsInput| {
            Ok(protocol::Mutation {
                request_id: input.request_id,
                branch: input.branch,
                expected: input.expected.0,
                change: protocol::Change::EdgeDelete(protocol::EdgeDelete {
                    relation: input.relation,
                    from: node(input.from),
                    to: node(input.to),
                }),
            })
        },
        |receipt| types::UnlinkRecordsOutput {
            branch: receipt.branch,
            version: receipt.version,
            request_id: receipt.request_id,
        },
    )
}

pub(super) fn create_branch(
    sender: &mut impl Sender,
    writer: Writer,
    call: &rpc::ToolCall,
) -> Result<ToolResult, EncodeError> {
    apply(
        sender,
        writer,
        call,
        |input: types::CreateBranchInput| {
            Ok(protocol::Mutation {
                request_id: input.request_id,
                branch: input.branch,
                expected: input.expected.0,
                change: protocol::Change::Fork(protocol::Fork {
                    name: input.name,
                    from: input.from.0,
                }),
            })
        },
        |receipt| types::CreateBranchOutput {
            branch: receipt.branch,
            version: receipt.version,
            request_id: receipt.request_id,
        },
    )
}

pub(super) fn merge_branch(
    sender: &mut impl Sender,
    writer: Writer,
    call: &rpc::ToolCall,
) -> Result<ToolResult, EncodeError> {
    apply(
        sender,
        writer,
        call,
        |input: types::MergeBranchInput| {
            Ok(protocol::Mutation {
                request_id: input.request_id,
                branch: input.branch,
                expected: input.expected.0,
                change: protocol::Change::Merge(protocol::BranchMerge {
                    from: input.from,
                    reviewed: input.reviewed.0,
                }),
            })
        },
        |receipt| types::MergeBranchOutput {
            branch: receipt.branch,
            version: receipt.version,
            request_id: receipt.request_id,
        },
    )
}

fn node(record: types::RecordRef) -> protocol::Node {
    protocol::Node {
        collection: record.collection,
        document_id: record.document_id,
    }
}