kaynine_runtime/

service.rs

//! AgentRuntime: the host-facing facade over per-session actors (SPEC §3.2,
//! §12). Each session has at most one actor task holding the writer lease,
//! which serializes all authoritative writes for that session.

use crate::actor::{self, ActorCommand, ActorHandle};
use crate::run;
use kaynine_core::budget::BudgetPolicy;
use kaynine_core::compaction::{CompactionConfig, CompactionModelSelector};
use kaynine_core::error::KaynineError;
use kaynine_core::event::{EventEnvelope, RealtimeEvent};
use kaynine_core::ids::{BranchId, ModelId, RunId, SessionId};
use kaynine_core::message::ContentBlock;
use kaynine_core::policy::Policy;
use kaynine_core::provider::{
    CredentialProvider, ModelCapabilities, ModelProvider, ReasoningLevel, TokenCounter,
};
use kaynine_core::store::{
    BranchRecord, CreateSessionRequest, EntryRecord, RunState, SessionRecord, SessionStore,
    SteerRecord,
};
use kaynine_core::tool::Tool;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::sync::{broadcast, mpsc, oneshot};

pub struct AgentRuntime {
    store: Arc<dyn SessionStore>,
    actors: Arc<Mutex<HashMap<SessionId, ActorHandle>>>,
    actor_seq: AtomicU64,
    shutting_down: AtomicBool,
}

#[derive(Clone)]
pub struct StartRunRequest {
    pub command_id: String,
    pub session_id: SessionId,
    pub branch_id: BranchId,
    pub content: Vec<ContentBlock>,
    pub model_override: Option<ModelId>,
    pub reasoning_override: Option<ReasoningLevel>,
    pub capabilities: ModelCapabilities,
    pub system_prompt: String,
    pub provider: Arc<dyn ModelProvider>,
    pub token_counter: Arc<dyn TokenCounter>,
    pub credentials: Arc<dyn CredentialProvider>,
    pub tools: Vec<Arc<dyn Tool>>,
    pub budget: BudgetPolicy,
    pub max_turns: Option<u32>,
    /// Policy gate evaluated before every tool execution (SPEC §8). Defaults
    /// to `AllowAllPolicy` for trusted hosts.
    pub policy: Arc<dyn Policy>,
    /// When set, the run gets an InteractiveApprovalHandler with this
    /// per-decision timeout; Ask decisions wait for `resolve_approval`.
    /// None → approval mode DisabledFailClosed: Ask fails closed.
    pub approval_timeout: Option<Duration>,
    /// When set, the loop composes the system prompt through this composer
    /// at every turn (SPEC §9); None falls back to `system_prompt`.
    pub prompt: Option<Arc<kaynine_core::prompt::PromptComposer>>,
    /// Compaction thresholds (SPEC §10); None disables compaction.
    pub compaction: Option<CompactionConfig>,
    /// Model selector for summary generation; None → current model.
    pub compaction_selector: Option<Arc<dyn CompactionModelSelector>>,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct RunAccepted {
    pub run_id: RunId,
    pub user_entry_id: kaynine_core::ids::EntryId,
    pub revision: u64,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct CancelRequest {
    pub command_id: String,
    pub session_id: SessionId,
    pub run_id: RunId,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum CancelOutcome {
    Cancelled { revision: u64 },
    AlreadyTerminal { state: RunState },
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct SessionSnapshot {
    pub session: SessionRecord,
    pub branches: Vec<BranchRecord>,
    pub chains: HashMap<BranchId, Vec<EntryRecord>>,
    /// Only reflects a live actor's run. A `Running` row left behind by a
    /// crashed process is surfaced after recovery-on-next-actor, not here.
    pub active_run: Option<ActiveRunInfo>,
    pub current_revision: u64,
    pub last_run_seq: Option<u64>,
    pub unapplied_steers: Vec<SteerRecord>,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ActiveRunInfo {
    pub run_id: RunId,
    pub branch_id: BranchId,
    pub model: ModelId,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum ReleaseOutcome {
    Released,
    RunAlreadyActive,
    NotFound,
}

#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct ShutdownReport {
    pub cancelled_runs: Vec<RunId>,
    pub interrupted_runs: Vec<RunId>,
    pub released_sessions: usize,
}

/// Items a subscription yields: the authoritative `Snapshot` first, then
/// live `Event`s. `ResyncRequired` appears when the consumer lagged past the
/// broadcast capacity; the following `next()` returns a fresh `Snapshot`.
/// Events between the initial registration and the snapshot (or spanning a
/// resync) can be delivered both in a snapshot and as events; consumers
/// dedupe by (revision, run_seq) (SPEC §5.3 allows this interpretation).
pub enum SubscriptionItem {
    Snapshot(Box<SessionSnapshot>),
    Event(EventEnvelope<RealtimeEvent>),
    ResyncRequired,
}

pub struct SessionSubscription {
    store: Arc<dyn SessionStore>,
    session_id: SessionId,
    receiver: broadcast::Receiver<EventEnvelope<RealtimeEvent>>,
    pending_snapshot: Option<Box<SessionSnapshot>>,
    resynced: bool,
}

impl SessionSubscription {
    pub async fn next(&mut self) -> Option<SubscriptionItem> {
        if let Some(snapshot) = self.pending_snapshot.take() {
            return Some(SubscriptionItem::Snapshot(snapshot));
        }
        if self.resynced {
            self.resynced = false;
            let snapshot =
                run::build_snapshot(self.store.as_ref(), &self.session_id, None, None, None)
                    .await
                    .ok()?;
            return Some(SubscriptionItem::Snapshot(Box::new(snapshot)));
        }
        match self.receiver.recv().await {
            Ok(envelope) => Some(SubscriptionItem::Event(envelope)),
            Err(broadcast::error::RecvError::Lagged(_)) => {
                self.resynced = true;
                Some(SubscriptionItem::ResyncRequired)
            }
            Err(broadcast::error::RecvError::Closed) => None,
        }
    }
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct UpdateSessionRequest {
    pub command_id: Option<String>,
    pub session_id: SessionId,
    pub default_model: Option<ModelId>,
    pub reasoning: Option<ReasoningLevel>,
    pub metadata: Option<serde_json::Value>,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct SteerRequest {
    pub command_id: String,
    pub session_id: SessionId,
    pub run_id: RunId,
    pub content: String,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct SteerAccepted {
    pub steer_id: String,
    pub revision: u64,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ApprovalResolution {
    pub command_id: String,
    pub session_id: SessionId,
    pub run_id: RunId,
    pub call_id: kaynine_core::ids::ToolCallId,
    pub approved: bool,
}

/// Result of submitting an approval decision (SPEC §8.3): Delivered reached a
/// live waiter; Expired arrived after the deadline lapsed; NotFound is an
/// unknown or already-resolved call.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum ApprovalOutcome {
    Delivered,
    Expired,
    NotFound,
}

impl AgentRuntime {
    pub fn new(store: Arc<dyn SessionStore>) -> Self {
        Self {
            store,
            actors: Arc::new(Mutex::new(HashMap::new())),
            actor_seq: AtomicU64::new(0),
            shutting_down: AtomicBool::new(false),
        }
    }

    pub async fn create_session(
        &self,
        request: CreateSessionRequest,
    ) -> Result<SessionRecord, KaynineError> {
        self.store.create_session(request).await
    }

    pub async fn list_sessions(&self) -> Result<Vec<SessionRecord>, KaynineError> {
        self.store.list_sessions().await
    }

    pub async fn list_branches(
        &self,
        session_id: &SessionId,
    ) -> Result<Vec<BranchRecord>, KaynineError> {
        self.store.list_branches(session_id).await
    }

    pub async fn get_snapshot(
        &self,
        session_id: &SessionId,
    ) -> Result<SessionSnapshot, KaynineError> {
        let handle = {
            let actors = self.actors.lock().expect("actor registry mutex poisoned");
            actors.get(session_id).cloned()
        };
        if let Some(handle) = handle {
            let (tx, rx) = oneshot::channel();
            if handle
                .tx
                .send(ActorCommand::Snapshot { reply: tx })
                .await
                .is_ok()
            {
                if let Ok(result) = rx.await {
                    return result;
                }
            }
            // The actor died mid-command; drop the stale registration and
            // fall through to reading store truth directly.
            self.remove_actor(session_id, &handle);
        }
        run::build_snapshot(self.store.as_ref(), session_id, None, None, None).await
    }

    pub async fn start_run(&self, request: StartRunRequest) -> Result<RunAccepted, KaynineError> {
        self.check_shutting_down()?;
        let session_id = request.session_id.clone();
        let request = Box::new(request);
        self.with_actor(&session_id, move |reply| ActorCommand::StartRun {
            request: request.clone(),
            reply,
        })
        .await
    }

    pub async fn cancel(&self, request: CancelRequest) -> Result<CancelOutcome, KaynineError> {
        self.check_shutting_down()?;
        let session_id = request.session_id.clone();
        self.with_actor(&session_id, move |reply| ActorCommand::Cancel {
            request: request.clone(),
            reply,
        })
        .await
    }

    /// Watches a session. Spawns (or reuses) the session actor, so if another
    /// process holds the writer lease this returns `SessionBusy` — v0.1
    /// targets single-process desktop use, where multi-process watching can
    /// read the store directly instead of via a live feed.
    pub async fn watch_session(
        &self,
        session_id: &SessionId,
    ) -> Result<SessionSubscription, KaynineError> {
        // Register the receiver BEFORE requesting the snapshot so no event
        // between the two is missed (duplicates in the other direction are
        // possible; consumers dedupe by revision/run_seq).
        let receiver = self
            .with_actor(session_id, |reply| ActorCommand::Subscribe { reply })
            .await?;
        let snapshot = self
            .with_actor(session_id, |reply| ActorCommand::Snapshot { reply })
            .await?;
        Ok(SessionSubscription {
            store: self.store.clone(),
            session_id: session_id.clone(),
            receiver,
            pending_snapshot: Some(Box::new(snapshot)),
            resynced: false,
        })
    }

    pub async fn steer(&self, request: SteerRequest) -> Result<SteerAccepted, KaynineError> {
        self.check_shutting_down()?;
        let session_id = request.session_id.clone();
        self.with_actor(&session_id, move |reply| ActorCommand::Steer {
            request: request.clone(),
            reply,
        })
        .await
    }

    pub async fn resolve_approval(
        &self,
        request: ApprovalResolution,
    ) -> Result<ApprovalOutcome, KaynineError> {
        self.check_shutting_down()?;
        let session_id = request.session_id.clone();
        self.with_actor(&session_id, move |reply| ActorCommand::ResolveApproval {
            request: request.clone(),
            reply,
        })
        .await
    }

    pub async fn update_session(
        &self,
        request: UpdateSessionRequest,
    ) -> Result<SessionRecord, KaynineError> {
        self.check_shutting_down()?;
        let session_id = request.session_id.clone();
        self.with_actor(&session_id, move |reply| ActorCommand::UpdateSession {
            request: request.clone(),
            reply,
        })
        .await
    }

    pub async fn release_session(
        &self,
        session_id: &SessionId,
    ) -> Result<ReleaseOutcome, KaynineError> {
        let handle = {
            let actors = self.actors.lock().expect("actor registry mutex poisoned");
            actors.get(session_id).cloned()
        };
        let Some(handle) = handle else {
            return Ok(ReleaseOutcome::NotFound);
        };
        let (tx, rx) = oneshot::channel();
        if handle
            .tx
            .send(ActorCommand::Release { reply: tx })
            .await
            .is_err()
        {
            self.remove_actor(session_id, &handle);
            return Ok(ReleaseOutcome::NotFound);
        }
        match rx.await {
            Ok(result) => result,
            Err(_) => {
                self.remove_actor(session_id, &handle);
                Ok(ReleaseOutcome::NotFound)
            }
        }
    }

    pub async fn shutdown(&self, grace: Duration) -> Result<ShutdownReport, KaynineError> {
        self.shutting_down.store(true, Ordering::SeqCst);
        let handles: Vec<(SessionId, ActorHandle)> = {
            let actors = self.actors.lock().expect("actor registry mutex poisoned");
            actors
                .iter()
                .map(|(id, h)| (id.clone(), h.clone()))
                .collect()
        };
        let mut report = ShutdownReport::default();
        for (session_id, handle) in handles {
            let (tx, rx) = oneshot::channel();
            if handle
                .tx
                .send(ActorCommand::Shutdown { grace, reply: tx })
                .await
                .is_err()
            {
                self.remove_actor(&session_id, &handle);
                continue;
            }
            if let Ok(Ok((cancelled, interrupted))) = rx.await {
                report.cancelled_runs.extend(cancelled);
                report.interrupted_runs.extend(interrupted);
                report.released_sessions += 1;
            }
            self.remove_actor(&session_id, &handle);
        }
        Ok(report)
    }

    fn check_shutting_down(&self) -> Result<(), KaynineError> {
        if self.shutting_down.load(Ordering::SeqCst) {
            return Err(KaynineError::InvalidRequest);
        }
        Ok(())
    }

    /// Sends a command to the session's actor, spawning one if needed. If the
    /// actor task already exited (e.g. its lease acquisition failed and it
    /// replied SessionBusy to queued commands before this send), the stale
    /// registration is dropped and a fresh actor is tried once.
    async fn with_actor<T>(
        &self,
        session_id: &SessionId,
        build: impl Fn(oneshot::Sender<Result<T, KaynineError>>) -> ActorCommand,
    ) -> Result<T, KaynineError> {
        let mut handle = self.get_or_spawn_actor(session_id)?;
        for attempt in 0..2 {
            let (tx, rx) = oneshot::channel();
            if handle.tx.send(build(tx)).await.is_ok() {
                return match rx.await {
                    Ok(result) => result,
                    Err(_) => {
                        if attempt == 0 {
                            self.remove_actor(session_id, &handle);
                            handle = self.get_or_spawn_actor(session_id)?;
                            continue;
                        }
                        Err(KaynineError::Internal)
                    }
                };
            }
            if attempt == 0 {
                self.remove_actor(session_id, &handle);
                handle = self.get_or_spawn_actor(session_id)?;
                continue;
            }
            return Err(KaynineError::Internal);
        }
        Err(KaynineError::Internal)
    }

    fn get_or_spawn_actor(&self, session_id: &SessionId) -> Result<ActorHandle, KaynineError> {
        self.check_shutting_down()?;
        let mut actors = self.actors.lock().expect("actor registry mutex poisoned");
        if let Some(handle) = actors.get(session_id) {
            return Ok(handle.clone());
        }
        let (tx, rx) = mpsc::channel(64);
        let handle = ActorHandle {
            actor_id: self.actor_seq.fetch_add(1, Ordering::SeqCst),
            tx,
        };
        actors.insert(session_id.clone(), handle.clone());
        drop(actors);
        actor::spawn(
            self.store.clone(),
            self.actors.clone(),
            session_id.clone(),
            handle.clone(),
            rx,
        );
        Ok(handle)
    }

    fn remove_actor(&self, session_id: &SessionId, handle: &ActorHandle) {
        let mut actors = self.actors.lock().expect("actor registry mutex poisoned");
        if actors
            .get(session_id)
            .is_some_and(|current| current.actor_id == handle.actor_id)
        {
            actors.remove(session_id);
        }
    }
}