kcode-commit-session 0.1.2

#![forbid(unsafe_code)]

use std::{
    collections::{BTreeMap, BTreeSet},
    fmt,
    path::Path,
    str::FromStr,
};

use chrono::{DateTime, Utc};
use kcode_kweb_db::{
    Error as KwebError, KwebDb, NodeData, NodeId, ObjectId, Owner, Provenance, TransactionId,
};
use rusqlite::{Connection, OptionalExtension, params};
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use sha2::{Digest, Sha256};

const DIGEST_VERSION: i64 = 2;
const RECEIPT_TABLE: &str = "kmap_session_commit_receipts";

/// Everything needed to materialize one sealed session.
///
/// Object payloads must already be in the permanent format expected by their
/// readers. Keys in `objects` and `creates` are canonical `pending:N` IDs.
#[derive(Clone, Debug)]
pub struct CommitRequest {
    pub idempotency_key: String,
    pub author: String,
    pub source_created_at: DateTime<Utc>,
    pub archive: Vec<u8>,
    pub objects: BTreeMap<String, Vec<u8>>,
    pub creates: BTreeMap<String, PlannedNode>,
    pub updates: BTreeMap<NodeId, PlannedNode>,
}

/// A node whose pending references will be resolved during commit.
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct PlannedNode {
    pub short_name: String,
    pub short_description: String,
    pub long_description: String,
    pub owner: String,
    #[serde(default)]
    pub fixed_connections: Vec<String>,
    #[serde(default)]
    pub recent_connections: Vec<String>,
    #[serde(default)]
    pub objects: Vec<String>,
    /// The serialized name preserves existing Kennedy checkpoints.
    #[serde(
        default,
        rename = "includeSessionObject",
        alias = "attachSessionArchive"
    )]
    pub attach_session_archive: bool,
}

/// The stable identifiers allocated for a committed session.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct CommitReceipt {
    pub transaction_id: Option<TransactionId>,
    pub session_object_id: ObjectId,
    pub node_ids: BTreeMap<String, NodeId>,
    pub object_ids: BTreeMap<String, ObjectId>,
}

#[derive(Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
struct WireReceipt {
    transaction_id: Option<String>,
    session_object_id: String,
    node_ids: BTreeMap<String, String>,
    object_ids: BTreeMap<String, String>,
}

impl Serialize for CommitReceipt {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        WireReceipt {
            transaction_id: self.transaction_id.map(|id| id.to_string()),
            session_object_id: self.session_object_id.to_string(),
            node_ids: self
                .node_ids
                .iter()
                .map(|(pending, id)| (pending.clone(), id.to_string()))
                .collect(),
            object_ids: self
                .object_ids
                .iter()
                .map(|(pending, id)| (pending.clone(), id.to_string()))
                .collect(),
        }
        .serialize(serializer)
    }
}

impl<'de> Deserialize<'de> for CommitReceipt {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let wire = WireReceipt::deserialize(deserializer)?;
        Ok(Self {
            transaction_id: wire
                .transaction_id
                .map(|value| TransactionId::from_str(&value).map_err(serde::de::Error::custom))
                .transpose()?,
            session_object_id: ObjectId::from_str(&wire.session_object_id)
                .map_err(serde::de::Error::custom)?,
            node_ids: parse_id_map(wire.node_ids)?,
            object_ids: parse_id_map(wire.object_ids)?,
        })
    }
}

fn parse_id_map<T, E>(values: BTreeMap<String, String>) -> Result<BTreeMap<String, T>, E>
where
    T: FromStr,
    T::Err: fmt::Display,
    E: serde::de::Error,
{
    values
        .into_iter()
        .map(|(key, value)| T::from_str(&value).map(|id| (key, id)).map_err(E::custom))
        .collect()
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[non_exhaustive]
pub enum ErrorKind {
    InvalidInput,
    NotFound,
    Conflict,
    Internal,
}

#[derive(Debug)]
pub struct Error {
    kind: ErrorKind,
    message: String,
}

impl Error {
    pub fn kind(&self) -> ErrorKind {
        self.kind
    }

    fn invalid(message: impl Into<String>) -> Self {
        Self {
            kind: ErrorKind::InvalidInput,
            message: message.into(),
        }
    }

    fn conflict(message: impl Into<String>) -> Self {
        Self {
            kind: ErrorKind::Conflict,
            message: message.into(),
        }
    }

    fn internal(error: impl fmt::Display) -> Self {
        Self {
            kind: ErrorKind::Internal,
            message: error.to_string(),
        }
    }
}

impl fmt::Display for Error {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(&self.message)
    }
}

impl std::error::Error for Error {}

impl From<KwebError> for Error {
    fn from(error: KwebError) -> Self {
        let kind = match error {
            KwebError::InvalidInput(_) | KwebError::InvalidTransaction(_) => {
                ErrorKind::InvalidInput
            }
            KwebError::NotFound(_) => ErrorKind::NotFound,
            _ => ErrorKind::Internal,
        };
        Self {
            kind,
            message: error.to_string(),
        }
    }
}

/// Atomically materializes a sealed session into one Kweb transaction.
///
/// Repeating an identical request returns its receipt. Reusing the
/// idempotency key for a different request returns `ErrorKind::Conflict`.
/// The caller must serialize this operation with other writers of the receipt
/// database; Kweb serializes its own transaction writer internally.
pub fn commit_session(
    database: &KwebDb,
    receipt_database: &Path,
    request: CommitRequest,
) -> Result<CommitReceipt, Error> {
    validate(&request)?;
    let request_digest = request_digest(&request);
    let receipts = open_receipts(receipt_database)?;
    if let Some(receipt) = recover_or_replay(database, &receipts, &request, request_digest)? {
        return Ok(receipt);
    }

    receipts
        .execute(
            "INSERT INTO kmap_session_commit_receipts(
                 session_id,request_sha256,digest_version,prepared_json,result_json,
                 started_at,committed_at
             ) VALUES(?1,?2,?3,NULL,NULL,?4,NULL)",
            params![
                &request.idempotency_key,
                request_digest.as_slice(),
                DIGEST_VERSION,
                Utc::now().to_rfc3339(),
            ],
        )
        .map_err(Error::internal)?;

    let mut transaction = database.start_transaction(Provenance {
        author: request.author.clone(),
        source: "kennedy-session".into(),
        source_created_at: request.source_created_at,
        data: format!("Kennedy session {}.", request.idempotency_key),
    })?;

    let mut object_ids = BTreeMap::new();
    for (pending, payload) in request.objects {
        object_ids.insert(pending, transaction.create_object(payload)?);
    }
    let archive = replace_pending_object_tokens(&request.archive, &object_ids);
    let session_object_id = transaction.create_object(archive)?;

    let mut node_ids = BTreeMap::new();
    for pending in request.creates.keys() {
        node_ids.insert(pending.clone(), transaction.reserve_node_id()?);
    }
    for (pending, data) in request.creates {
        let resolved = resolve_node(data, &node_ids, &object_ids, session_object_id)?;
        transaction.create_reserved_node(node_ids[&pending], resolved)?;
    }
    for (id, data) in request.updates {
        let resolved = resolve_node(data, &node_ids, &object_ids, session_object_id)?;
        transaction.update_node(id, resolved)?;
    }

    let mut receipt = CommitReceipt {
        transaction_id: None,
        session_object_id,
        node_ids,
        object_ids,
    };
    let prepared_json = serde_json::to_string(&receipt).map_err(Error::internal)?;
    require_one(
        receipts
            .execute(
                "UPDATE kmap_session_commit_receipts
                 SET prepared_json=?2
                 WHERE session_id=?1 AND prepared_json IS NULL AND result_json IS NULL",
                params![&request.idempotency_key, prepared_json],
            )
            .map_err(Error::internal)?,
        "session commit preparation receipt disappeared",
    )?;

    receipt.transaction_id = Some(transaction.finalize()?);
    let result_json = serde_json::to_string(&receipt).map_err(Error::internal)?;
    require_one(
        receipts
            .execute(
                "UPDATE kmap_session_commit_receipts
                 SET result_json=?2,committed_at=?3
                 WHERE session_id=?1 AND result_json IS NULL",
                params![
                    &request.idempotency_key,
                    result_json,
                    Utc::now().to_rfc3339(),
                ],
            )
            .map_err(Error::internal)?,
        "session commit receipt disappeared during mutation",
    )?;
    Ok(receipt)
}

fn open_receipts(path: &Path) -> Result<Connection, Error> {
    let connection = Connection::open(path).map_err(Error::internal)?;
    connection
        .execute_batch(
            "PRAGMA foreign_keys=ON;
             PRAGMA journal_mode=WAL;
             PRAGMA busy_timeout=15000;
             CREATE TABLE IF NOT EXISTS kmap_session_commit_receipts (
                 session_id TEXT PRIMARY KEY,
                 request_sha256 BLOB NOT NULL CHECK(length(request_sha256)=32),
                 digest_version INTEGER NOT NULL DEFAULT 2,
                 prepared_json TEXT,
                 result_json TEXT,
                 started_at TEXT NOT NULL,
                 committed_at TEXT,
                 CHECK((result_json IS NULL) = (committed_at IS NULL))
             );",
        )
        .map_err(Error::internal)?;
    let columns = {
        let mut statement = connection
            .prepare(&format!("PRAGMA table_info({RECEIPT_TABLE})"))
            .map_err(Error::internal)?;
        statement
            .query_map([], |row| row.get::<_, String>(1))
            .map_err(Error::internal)?
            .collect::<rusqlite::Result<BTreeSet<_>>>()
            .map_err(Error::internal)?
    };
    if !columns.contains("prepared_json") {
        connection
            .execute(
                "ALTER TABLE kmap_session_commit_receipts ADD COLUMN prepared_json TEXT",
                [],
            )
            .map_err(Error::internal)?;
    }
    if !columns.contains("digest_version") {
        connection
            .execute(
                "ALTER TABLE kmap_session_commit_receipts
                 ADD COLUMN digest_version INTEGER NOT NULL DEFAULT 1",
                [],
            )
            .map_err(Error::internal)?;
    }
    Ok(connection)
}

fn recover_or_replay(
    database: &KwebDb,
    receipts: &Connection,
    request: &CommitRequest,
    request_digest: [u8; 32],
) -> Result<Option<CommitReceipt>, Error> {
    let existing = receipts
        .query_row(
            "SELECT request_sha256,digest_version,prepared_json,result_json
             FROM kmap_session_commit_receipts WHERE session_id=?1",
            [&request.idempotency_key],
            |row| {
                Ok((
                    row.get::<_, Vec<u8>>(0)?,
                    row.get::<_, i64>(1)?,
                    row.get::<_, Option<String>>(2)?,
                    row.get::<_, Option<String>>(3)?,
                ))
            },
        )
        .optional()
        .map_err(Error::internal)?;
    let Some((stored_digest, digest_version, prepared, result)) = existing else {
        return Ok(None);
    };
    // Version 1 was Kennedy's serde-derived digest. Its private request types
    // no longer exist here, so the already-unique key is the compatibility
    // authority for those historical receipts.
    if digest_version == DIGEST_VERSION && stored_digest.as_slice() != request_digest {
        return Err(Error::conflict(
            "idempotency key was already used for a different session commit",
        ));
    }
    if let Some(result) = result {
        return serde_json::from_str(&result)
            .map(Some)
            .map_err(Error::internal);
    }
    if let Some(prepared) = prepared {
        let recovered: CommitReceipt = serde_json::from_str(&prepared).map_err(Error::internal)?;
        match database.get_object(recovered.session_object_id) {
            Ok(_) => {
                require_one(
                    receipts
                        .execute(
                            "UPDATE kmap_session_commit_receipts
                             SET result_json=prepared_json,committed_at=?2
                             WHERE session_id=?1 AND result_json IS NULL",
                            params![&request.idempotency_key, Utc::now().to_rfc3339()],
                        )
                        .map_err(Error::internal)?,
                    "session recovery receipt disappeared",
                )?;
                return Ok(Some(recovered));
            }
            Err(KwebError::NotFound(_)) => {}
            Err(error) => return Err(error.into()),
        }
    }
    receipts
        .execute(
            "DELETE FROM kmap_session_commit_receipts
             WHERE session_id=?1 AND result_json IS NULL",
            [&request.idempotency_key],
        )
        .map_err(Error::internal)?;
    Ok(None)
}

fn validate(request: &CommitRequest) -> Result<(), Error> {
    if request.idempotency_key.trim().is_empty() || request.idempotency_key.len() > 1024 {
        return Err(Error::invalid(
            "idempotency key must contain between 1 and 1024 bytes",
        ));
    }
    let mut pending_ids = BTreeSet::new();
    for pending in request.objects.keys().chain(request.creates.keys()) {
        validate_pending_id(pending)?;
        if !pending_ids.insert(pending) {
            return Err(Error::invalid(format!(
                "pending ID {pending} is used for both an object and a node"
            )));
        }
    }
    for node in request.creates.values().chain(request.updates.values()) {
        validate_node(node, &request.creates, &request.objects)?;
    }
    Ok(())
}

fn validate_pending_id(value: &str) -> Result<(), Error> {
    let number = value
        .strip_prefix("pending:")
        .and_then(|number| number.parse::<u64>().ok())
        .filter(|number| *number > 0);
    if number.is_none() || format!("pending:{}", number.unwrap_or_default()) != value {
        return Err(Error::invalid(format!(
            "{value:?} is not a canonical pending ID"
        )));
    }
    Ok(())
}

fn validate_node(
    node: &PlannedNode,
    creates: &BTreeMap<String, PlannedNode>,
    objects: &BTreeMap<String, Vec<u8>>,
) -> Result<(), Error> {
    if !matches!(node.owner.as_str(), "self" | "unowned") {
        validate_node_ref(&node.owner, creates)?;
    }
    for value in node
        .fixed_connections
        .iter()
        .chain(&node.recent_connections)
    {
        validate_node_ref(value, creates)?;
    }
    for value in &node.objects {
        if value.starts_with("pending:") {
            validate_pending_id(value)?;
            if !objects.contains_key(value) {
                return Err(Error::invalid(format!("unresolved pending object {value}")));
            }
        } else {
            ObjectId::from_str(value).map_err(Error::from)?;
        }
    }
    Ok(())
}

fn validate_node_ref(value: &str, creates: &BTreeMap<String, PlannedNode>) -> Result<(), Error> {
    if value.starts_with("pending:") {
        validate_pending_id(value)?;
        if !creates.contains_key(value) {
            return Err(Error::invalid(format!("unresolved pending node {value}")));
        }
        Ok(())
    } else {
        NodeId::from_str(value).map(|_| ()).map_err(Error::from)
    }
}

fn resolve_node(
    node: PlannedNode,
    node_ids: &BTreeMap<String, NodeId>,
    object_ids: &BTreeMap<String, ObjectId>,
    session_object_id: ObjectId,
) -> Result<NodeData, Error> {
    let resolve_node_id = |value: &str| {
        if value.starts_with("pending:") {
            node_ids
                .get(value)
                .copied()
                .ok_or_else(|| Error::invalid(format!("unresolved pending node {value}")))
        } else {
            NodeId::from_str(value).map_err(Error::from)
        }
    };
    let owner = match node.owner.as_str() {
        "unowned" => Owner::Unowned,
        "self" => Owner::SelfNode,
        value => Owner::Node(resolve_node_id(value)?),
    };
    let mut objects = node
        .objects
        .iter()
        .map(|value| {
            if value.starts_with("pending:") {
                object_ids
                    .get(value)
                    .copied()
                    .ok_or_else(|| Error::invalid(format!("unresolved pending object {value}")))
            } else {
                ObjectId::from_str(value).map_err(Error::from)
            }
        })
        .collect::<Result<Vec<_>, _>>()?;
    if node.attach_session_archive && !objects.contains(&session_object_id) {
        objects.push(session_object_id);
    }
    Ok(NodeData {
        short_name: replace_pending_object_tokens_in_text(&node.short_name, object_ids),
        short_description: replace_pending_object_tokens_in_text(
            &node.short_description,
            object_ids,
        ),
        long_description: replace_pending_object_tokens_in_text(&node.long_description, object_ids),
        owner,
        fixed_connections: node
            .fixed_connections
            .iter()
            .map(|value| resolve_node_id(value))
            .collect::<Result<_, _>>()?,
        recent_connections: node
            .recent_connections
            .iter()
            .map(|value| resolve_node_id(value))
            .collect::<Result<_, _>>()?,
        objects,
    })
}

fn replace_pending_object_tokens(bytes: &[u8], object_ids: &BTreeMap<String, ObjectId>) -> Vec<u8> {
    let Ok(text) = std::str::from_utf8(bytes) else {
        return bytes.to_vec();
    };
    replace_pending_object_tokens_in_text(text, object_ids).into_bytes()
}

fn replace_pending_object_tokens_in_text(
    text: &str,
    object_ids: &BTreeMap<String, ObjectId>,
) -> String {
    if object_ids.is_empty() || !text.contains("pending:") {
        return text.into();
    }
    let mut output = String::with_capacity(text.len());
    let mut cursor = 0;
    while let Some(relative) = text[cursor..].find("pending:") {
        let start = cursor + relative;
        let number_start = start + "pending:".len();
        let number_len = text[number_start..]
            .bytes()
            .take_while(u8::is_ascii_digit)
            .count();
        if number_len == 0 {
            output.push_str(&text[cursor..number_start]);
            cursor = number_start;
            continue;
        }
        let end = number_start + number_len;
        let token = &text[start..end];
        let left_boundary = start == 0
            || !text[..start]
                .chars()
                .next_back()
                .is_some_and(|character| character.is_ascii_alphanumeric() || character == '_');
        let right_boundary = text[end..]
            .chars()
            .next()
            .is_none_or(|character| !character.is_ascii_alphanumeric() && character != '_');
        if left_boundary
            && right_boundary
            && let Some(id) = object_ids.get(token)
        {
            output.push_str(&text[cursor..start]);
            output.push_str(&id.to_string());
            cursor = end;
            continue;
        }
        output.push_str(&text[cursor..end]);
        cursor = end;
    }
    output.push_str(&text[cursor..]);
    output
}

fn request_digest(request: &CommitRequest) -> [u8; 32] {
    let mut digest = Sha256::new();
    digest.update(b"kcode-commit-session request v2\0");
    hash_bytes(&mut digest, request.idempotency_key.as_bytes());
    hash_bytes(&mut digest, request.author.as_bytes());
    digest.update(request.source_created_at.timestamp().to_be_bytes());
    digest.update(
        request
            .source_created_at
            .timestamp_subsec_nanos()
            .to_be_bytes(),
    );
    hash_bytes(&mut digest, &request.archive);
    hash_u64(&mut digest, request.objects.len());
    for (pending, bytes) in &request.objects {
        hash_bytes(&mut digest, pending.as_bytes());
        hash_bytes(&mut digest, bytes);
    }
    hash_u64(&mut digest, request.creates.len());
    for (pending, node) in &request.creates {
        hash_bytes(&mut digest, pending.as_bytes());
        hash_node(&mut digest, node);
    }
    hash_u64(&mut digest, request.updates.len());
    for (id, node) in &request.updates {
        digest.update(id.to_bytes());
        hash_node(&mut digest, node);
    }
    digest.finalize().into()
}

fn hash_node(digest: &mut Sha256, node: &PlannedNode) {
    hash_bytes(digest, node.short_name.as_bytes());
    hash_bytes(digest, node.short_description.as_bytes());
    hash_bytes(digest, node.long_description.as_bytes());
    hash_bytes(digest, node.owner.as_bytes());
    hash_strings(digest, &node.fixed_connections);
    hash_strings(digest, &node.recent_connections);
    hash_strings(digest, &node.objects);
    digest.update([u8::from(node.attach_session_archive)]);
}

fn hash_strings(digest: &mut Sha256, values: &[String]) {
    hash_u64(digest, values.len());
    for value in values {
        hash_bytes(digest, value.as_bytes());
    }
}

fn hash_bytes(digest: &mut Sha256, bytes: &[u8]) {
    hash_u64(digest, bytes.len());
    digest.update(bytes);
}

fn hash_u64(digest: &mut Sha256, value: usize) {
    digest.update((value as u64).to_be_bytes());
}

fn require_one(updated: usize, message: &'static str) -> Result<(), Error> {
    if updated == 1 {
        Ok(())
    } else {
        Err(Error::internal(message))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn planned_node_keeps_checkpoint_field_compatible() {
        let node = PlannedNode {
            short_name: "name".into(),
            short_description: String::new(),
            long_description: String::new(),
            owner: "self".into(),
            fixed_connections: Vec::new(),
            recent_connections: Vec::new(),
            objects: Vec::new(),
            attach_session_archive: true,
        };
        let value = serde_json::to_value(&node).unwrap();
        assert_eq!(value["includeSessionObject"], true);
        assert_eq!(
            serde_json::from_value::<PlannedNode>(serde_json::json!({
                "shortName": "name",
                "shortDescription": "",
                "longDescription": "",
                "owner": "self",
                "attachSessionArchive": true
            }))
            .unwrap(),
            node
        );
    }
}