kaynine-core 0.1.0

use crate::error::KaynineError;
use crate::error::RunFailureReason;
use crate::ids::{BranchId, EntryId, ModelId, ProviderId, RunId, SessionId, ToolCallId};
use crate::message::{BinaryRef, FinishReason, Message, Usage};
use crate::provider::ReasoningLevel;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct EntryRecord {
    pub entry_id: EntryId,
    pub parent_entry_id: Option<EntryId>,
    pub branch_id: BranchId,
    pub message: Message,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct SessionRecord {
    pub session_id: SessionId,
    pub definition_id: String,
    pub definition_version: u32,
    pub default_model: ModelId,
    pub created_at_unix: i64,
    pub updated_at_unix: i64,
    pub current_revision: u64,
    pub metadata: serde_json::Value,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct BranchRecord {
    pub branch_id: BranchId,
    pub session_id: SessionId,
    pub head_entry_id: Option<EntryId>,
    pub needs_compaction: bool,
    pub created_at_unix: i64,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum RunState {
    Running,
    Completed,
    Cancelled,
    Failed,
    Interrupted,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct RunRecord {
    pub run_id: RunId,
    pub session_id: SessionId,
    pub branch_id: BranchId,
    pub state: RunState,
    pub model: ModelId,
    pub reasoning: ReasoningLevel,
    pub failure: Option<RunFailureReason>,
    pub started_at_unix: i64,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum ToolCallState {
    Planned,
    Executing,
    Succeeded,
    Failed,
    InDoubt,
    CancelledBySystem,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ToolCallRecord {
    pub call_id: ToolCallId,
    pub run_id: RunId,
    pub turn: u32,
    pub name: String,
    pub arguments: serde_json::Value,
    pub state: ToolCallState,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum SteerState {
    Queued,
    Applied { entry_id: EntryId },
    Unapplied,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct SteerRecord {
    pub steer_id: String,
    pub session_id: SessionId,
    pub content: String,
    pub state: SteerState,
    pub received_at_unix: i64,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct SummaryRecord {
    pub summary_id: String,
    pub branch_id: BranchId,
    pub covered_until_entry: EntryId,
    pub text: String,
    pub source_hash: String,
    pub provider: ProviderId,
    pub model: ModelId,
    pub prompt_version: String,
    pub usage: Usage,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "kind")]
pub enum AuthoritativeEvent {
    UserMessageAccepted {
        entry: EntryRecord,
    },
    RunStarted {
        run: RunRecord,
    },
    TurnStarted {
        run_id: RunId,
        turn: u32,
    },
    TurnCompleted {
        run_id: RunId,
        turn: u32,
        finish_reason: FinishReason,
    },
    AssistantMessagePersisted {
        entry: EntryRecord,
    },
    ToolCallPlanned {
        call: ToolCallRecord,
    },
    PermissionDecided {
        call_id: ToolCallId,
        decision: String,
    },
    ApprovalRequested {
        call_id: ToolCallId,
        deadline_unix: i64,
    },
    ApprovalResolved {
        call_id: ToolCallId,
        approved: bool,
    },
    ToolResultRecorded {
        entry: EntryRecord,
        call_id: ToolCallId,
        synthesized: bool,
    },
    SteerQueued {
        steer: SteerRecord,
    },
    SteerApplied {
        steer_id: String,
        entry: EntryRecord,
    },
    ModelChanged {
        model: ModelId,
        reasoning: ReasoningLevel,
    },
    SessionMetadataUpdated {
        metadata: serde_json::Value,
    },
    BranchForked {
        new_branch: BranchRecord,
        from_branch_id: BranchId,
        from_entry_id: Option<EntryId>,
    },
    SummaryCheckpoint {
        summary: SummaryRecord,
    },
    RunTerminal {
        run_id: RunId,
        state: RunState,
        failure: Option<RunFailureReason>,
    },
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct LeaseOwner {
    pub owner_id: String,
    pub generation: u64,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "outcome")]
pub enum AppendOutcome {
    Appended { new_revision: u64 },
    RevisionConflict { current_revision: u64 },
    NotLeaseHolder,
    LeaseExpired,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct CommandRecord {
    pub command_id: String,
    pub kind: String,
    pub request_hash: String,
    pub result: Option<serde_json::Value>,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "outcome")]
pub enum CommandOutcome {
    Accepted,
    Duplicate(CommandRecord),
    Conflict,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct CreateSessionRequest {
    pub command: Option<CommandRecord>,
    pub session_id: SessionId,
    pub definition_id: String,
    pub definition_version: u32,
    pub default_model: ModelId,
    pub metadata: serde_json::Value,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ForkRequest {
    pub command: Option<CommandRecord>,
    pub session_id: SessionId,
    pub from_branch_id: BranchId,
    pub from_entry_id: EntryId,
    pub new_branch_id: BranchId,
    pub expected_revision: u64,
    pub owner: LeaseOwner,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "outcome")]
pub enum ForkOutcome {
    Forked {
        branch: BranchRecord,
        new_revision: u64,
    },
    RevisionConflict {
        current_revision: u64,
    },
    NotLeaseHolder,
    NotAncestorOfHead,
    SessionBusy,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct RecoveryReport {
    pub session_id: SessionId,
    pub new_revision: u64,
    pub interrupted_runs: Vec<RunId>,
    pub expired_approvals: Vec<ToolCallId>,
    pub in_doubt_calls: Vec<ToolCallId>,
    pub cancelled_calls: Vec<ToolCallId>,
    pub synthesized_results: Vec<ToolCallId>,
    pub unapplied_steers: Vec<String>,
}

#[async_trait]
pub trait SessionStore: Send + Sync {
    async fn migrate(&self) -> Result<(), KaynineError>;
    async fn create_session(
        &self,
        request: CreateSessionRequest,
    ) -> Result<SessionRecord, KaynineError>;
    async fn list_sessions(&self) -> Result<Vec<SessionRecord>, KaynineError>;
    async fn get_session(
        &self,
        session_id: &SessionId,
    ) -> Result<Option<SessionRecord>, KaynineError>;
    async fn list_branches(
        &self,
        session_id: &SessionId,
    ) -> Result<Vec<BranchRecord>, KaynineError>;
    async fn load_chain(&self, head_entry_id: &EntryId) -> Result<Vec<EntryRecord>, KaynineError>;
    async fn list_tool_calls(&self, run_id: &RunId) -> Result<Vec<ToolCallRecord>, KaynineError>;
    async fn get_run(&self, run_id: &RunId) -> Result<Option<RunRecord>, KaynineError>;
    async fn list_steers(&self, session_id: &SessionId) -> Result<Vec<SteerRecord>, KaynineError>;
    /// Compaction checkpoints recorded for a branch, oldest first (SPEC
    /// §10). Used to lazily restore the working context on the next run.
    async fn list_summaries(
        &self,
        branch_id: &BranchId,
    ) -> Result<Vec<SummaryRecord>, KaynineError>;
    /// Reads the payload of every persisted policy grant in `namespace`
    /// (SPEC §8.3 grant layer). Grants are written by the product layer, not
    /// by session events. Stores without a grants table return empty.
    async fn list_grants(&self, _namespace: &str) -> Result<Vec<serde_json::Value>, KaynineError> {
        Ok(Vec::new())
    }
    async fn append_events(
        &self,
        session_id: &SessionId,
        expected_revision: u64,
        owner: &LeaseOwner,
        events: Vec<AuthoritativeEvent>,
        blobs: Vec<(BinaryRef, Vec<u8>)>,
    ) -> Result<AppendOutcome, KaynineError>;
    async fn fork_branch(&self, request: ForkRequest) -> Result<ForkOutcome, KaynineError>;
    async fn record_command(&self, command: CommandRecord) -> Result<CommandOutcome, KaynineError>;
    async fn complete_command(
        &self,
        command_id: &str,
        result: serde_json::Value,
    ) -> Result<(), KaynineError>;
    async fn load_blob(&self, sha256: &str) -> Result<Option<Vec<u8>>, KaynineError>;
    async fn acquire_lease(
        &self,
        session_id: &SessionId,
        owner_id: &str,
        ttl_secs: i64,
    ) -> Result<LeaseOwner, KaynineError>;
    async fn renew_lease(
        &self,
        session_id: &SessionId,
        owner: &LeaseOwner,
        ttl_secs: i64,
    ) -> Result<bool, KaynineError>;
    async fn release_lease(
        &self,
        session_id: &SessionId,
        owner: &LeaseOwner,
    ) -> Result<(), KaynineError>;
    async fn recover_session(
        &self,
        session_id: &SessionId,
        owner: &LeaseOwner,
    ) -> Result<RecoveryReport, KaynineError>;
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;
    fn entry_record() -> EntryRecord {
        EntryRecord {
            entry_id: EntryId::from("entry-1"),
            parent_entry_id: Some(EntryId::from("entry-0")),
            branch_id: BranchId::from("branch-1"),
            message: Message::User {
                blocks: vec![crate::message::ContentBlock::Text {
                    text: "hello".to_string(),
                }],
            },
        }
    }

    fn run_record() -> RunRecord {
        RunRecord {
            run_id: RunId::from("run-1"),
            session_id: SessionId::from("session-1"),
            branch_id: BranchId::from("branch-1"),
            state: RunState::Running,
            model: ModelId::from("model-1"),
            reasoning: ReasoningLevel::Medium,
            failure: None,
            started_at_unix: 1,
        }
    }

    fn steer_record(state: SteerState) -> SteerRecord {
        SteerRecord {
            steer_id: "steer-1".to_string(),
            session_id: SessionId::from("session-1"),
            content: "steer".to_string(),
            state,
            received_at_unix: 2,
        }
    }

    fn summary_record() -> SummaryRecord {
        SummaryRecord {
            summary_id: "summary-1".to_string(),
            branch_id: BranchId::from("branch-1"),
            covered_until_entry: EntryId::from("entry-1"),
            text: "summary".to_string(),
            source_hash: "abc123".to_string(),
            provider: ProviderId::from("anthropic"),
            model: ModelId::from("model-1"),
            prompt_version: "v1".to_string(),
            usage: Usage {
                input_tokens: 10,
                output_tokens: 20,
                cached_input_tokens: Some(5),
            },
        }
    }

    fn branch_record() -> BranchRecord {
        BranchRecord {
            branch_id: BranchId::from("branch-2"),
            session_id: SessionId::from("session-1"),
            head_entry_id: Some(EntryId::from("entry-1")),
            needs_compaction: false,
            created_at_unix: 3,
        }
    }

    fn tool_call_record() -> ToolCallRecord {
        ToolCallRecord {
            call_id: ToolCallId::from("call-1"),
            run_id: RunId::from("run-1"),
            turn: 1,
            name: "bash".to_string(),
            arguments: json!({}),
            state: ToolCallState::Planned,
        }
    }

    fn roundtrip<T>(value: &T) -> T
    where
        T: Serialize + serde::de::DeserializeOwned,
    {
        let encoded = serde_json::to_string(value).unwrap();
        serde_json::from_str(&encoded).unwrap()
    }

    #[test]
    fn authoritative_events_roundtrip() {
        let events = vec![
            AuthoritativeEvent::UserMessageAccepted {
                entry: entry_record(),
            },
            AuthoritativeEvent::RunStarted { run: run_record() },
            AuthoritativeEvent::TurnStarted {
                run_id: RunId::from("run-1"),
                turn: 1,
            },
            AuthoritativeEvent::TurnCompleted {
                run_id: RunId::from("run-1"),
                turn: 1,
                finish_reason: FinishReason::Stop,
            },
            AuthoritativeEvent::AssistantMessagePersisted {
                entry: entry_record(),
            },
            AuthoritativeEvent::ToolCallPlanned {
                call: tool_call_record(),
            },
            AuthoritativeEvent::PermissionDecided {
                call_id: ToolCallId::from("call-1"),
                decision: "allow".to_string(),
            },
            AuthoritativeEvent::ApprovalRequested {
                call_id: ToolCallId::from("call-1"),
                deadline_unix: 100,
            },
            AuthoritativeEvent::ApprovalResolved {
                call_id: ToolCallId::from("call-1"),
                approved: true,
            },
            AuthoritativeEvent::ToolResultRecorded {
                entry: entry_record(),
                call_id: ToolCallId::from("call-1"),
                synthesized: true,
            },
            AuthoritativeEvent::SteerQueued {
                steer: steer_record(SteerState::Queued),
            },
            AuthoritativeEvent::SteerApplied {
                steer_id: "steer-1".to_string(),
                entry: entry_record(),
            },
            AuthoritativeEvent::ModelChanged {
                model: ModelId::from("model-2"),
                reasoning: ReasoningLevel::High,
            },
            AuthoritativeEvent::SessionMetadataUpdated {
                metadata: json!({"k": "v"}),
            },
            AuthoritativeEvent::BranchForked {
                new_branch: branch_record(),
                from_branch_id: BranchId::from("branch-1"),
                from_entry_id: Some(EntryId::from("entry-1")),
            },
            AuthoritativeEvent::SummaryCheckpoint {
                summary: summary_record(),
            },
            AuthoritativeEvent::RunTerminal {
                run_id: RunId::from("run-1"),
                state: RunState::Failed,
                failure: Some(RunFailureReason::Internal),
            },
        ];
        for event in &events {
            assert_eq!(&roundtrip(event), event);
        }
        assert_eq!(events.len(), 17);
    }

    #[test]
    fn records_roundtrip() {
        let session = SessionRecord {
            session_id: SessionId::from("session-1"),
            definition_id: "defs/agent".to_string(),
            definition_version: 1,
            default_model: ModelId::from("model-1"),
            created_at_unix: 1,
            updated_at_unix: 2,
            current_revision: 3,
            metadata: json!({}),
        };
        assert_eq!(roundtrip(&session), session);

        let branch = branch_record();
        assert_eq!(roundtrip(&branch), branch);

        let run = run_record();
        assert_eq!(roundtrip(&run), run);

        let call = tool_call_record();
        assert_eq!(roundtrip(&call), call);

        let steer_states = vec![
            SteerState::Queued,
            SteerState::Applied {
                entry_id: EntryId::from("entry-1"),
            },
            SteerState::Unapplied,
        ];
        for state in steer_states {
            let steer = steer_record(state);
            assert_eq!(roundtrip(&steer), steer);
        }

        let summary = summary_record();
        assert_eq!(roundtrip(&summary), summary);
    }

    #[test]
    fn outcomes_roundtrip() {
        let appends = vec![
            AppendOutcome::Appended { new_revision: 2 },
            AppendOutcome::RevisionConflict {
                current_revision: 1,
            },
            AppendOutcome::NotLeaseHolder,
            AppendOutcome::LeaseExpired,
        ];
        for outcome in &appends {
            assert_eq!(&roundtrip(outcome), outcome);
        }

        let command = CommandRecord {
            command_id: "cmd-1".to_string(),
            kind: "create_session".to_string(),
            request_hash: "hash".to_string(),
            result: Some(json!({"ok": true})),
        };
        let commands = vec![
            CommandOutcome::Accepted,
            CommandOutcome::Duplicate(command),
            CommandOutcome::Conflict,
        ];
        for outcome in &commands {
            assert_eq!(&roundtrip(outcome), outcome);
        }

        let forks = vec![
            ForkOutcome::Forked {
                branch: branch_record(),
                new_revision: 5,
            },
            ForkOutcome::RevisionConflict {
                current_revision: 4,
            },
            ForkOutcome::NotLeaseHolder,
            ForkOutcome::NotAncestorOfHead,
            ForkOutcome::SessionBusy,
        ];
        for outcome in &forks {
            assert_eq!(&roundtrip(outcome), outcome);
        }
    }
}