kaynine_runtime/

actor.rs

//! Per-session actor: the single writer for one session. Holds the writer
//! lease, keeps the in-memory current revision, tracks the active run, and
//! owns the broadcast feed for the session.

use crate::approval::ApprovalShared;
use crate::run::{self, build_snapshot, SteerQueue};
use crate::service::{
    ActiveRunInfo, ApprovalOutcome, ApprovalResolution, CancelOutcome, CancelRequest,
    ReleaseOutcome, RunAccepted, SessionSnapshot, StartRunRequest, SteerAccepted, SteerRequest,
    UpdateSessionRequest,
};
use kaynine_core::error::KaynineError;
use kaynine_core::event::{EventEnvelope, RealtimeEvent};
use kaynine_core::ids::{ModelId, RunId, SessionId};
use kaynine_core::store::{LeaseOwner, SessionRecord, SessionStore};
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::sync::{broadcast, mpsc, oneshot};
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;

const LEASE_TTL_SECS: i64 = 60;
const RENEW_INTERVAL_SECS: u64 = 20;
const BROADCAST_CAPACITY: usize = 256;

/// Per-actor shutdown result: (cancelled run ids, interrupted run ids).
pub(crate) type ActorShutdownOutcome = (Vec<RunId>, Vec<RunId>);

pub(crate) enum ActorCommand {
    StartRun {
        request: Box<StartRunRequest>,
        reply: oneshot::Sender<Result<RunAccepted, KaynineError>>,
    },
    Cancel {
        request: CancelRequest,
        reply: oneshot::Sender<Result<CancelOutcome, KaynineError>>,
    },
    Snapshot {
        reply: oneshot::Sender<Result<SessionSnapshot, KaynineError>>,
    },
    UpdateSession {
        request: UpdateSessionRequest,
        reply: oneshot::Sender<Result<SessionRecord, KaynineError>>,
    },
    Steer {
        request: SteerRequest,
        reply: oneshot::Sender<Result<SteerAccepted, KaynineError>>,
    },
    ResolveApproval {
        request: ApprovalResolution,
        reply: oneshot::Sender<Result<ApprovalOutcome, KaynineError>>,
    },
    Shutdown {
        grace: Duration,
        reply: oneshot::Sender<Result<ActorShutdownOutcome, KaynineError>>,
    },
    Release {
        reply: oneshot::Sender<Result<ReleaseOutcome, KaynineError>>,
    },
    Subscribe {
        reply: oneshot::Sender<
            Result<broadcast::Receiver<EventEnvelope<RealtimeEvent>>, KaynineError>,
        >,
    },
}

#[derive(Clone)]
pub(crate) struct ActorHandle {
    pub(crate) actor_id: u64,
    pub(crate) tx: mpsc::Sender<ActorCommand>,
}

pub(crate) type ActorRegistry = Arc<Mutex<HashMap<SessionId, ActorHandle>>>;

pub(crate) struct ActiveRun {
    pub(crate) run_id: RunId,
    pub(crate) branch_id: kaynine_core::ids::BranchId,
    pub(crate) model: ModelId,
    pub(crate) cancel: CancellationToken,
    pub(crate) join: JoinHandle<()>,
    /// Mirrors the loop's `cancel_grace`: once cancelled, the run task is
    /// guaranteed to append its terminal within this bound (+ epsilon), so
    /// the cancel handler can wait it out instead of detaching.
    pub(crate) cancel_grace: Duration,
    /// In-memory steer queue shared with the run task (producer here,
    /// drained by the run hooks when SteerApplied is persisted).
    pub(crate) steer_queue: Arc<SteerQueue>,
    /// Live approval waiters shared with the run task's
    /// InteractiveApprovalHandler; only populated when the run was started
    /// with an approval timeout.
    pub(crate) approval_shared: Arc<ApprovalShared>,
}

pub(crate) struct ActorState {
    pub(crate) session_id: SessionId,
    pub(crate) store: Arc<dyn SessionStore>,
    pub(crate) owner: LeaseOwner,
    /// Shared with the run task so hook appends and the actor agree on the
    /// expected revision (single writer, so no conflicts in practice).
    pub(crate) revision: Arc<tokio::sync::Mutex<u64>>,
    pub(crate) active: Option<ActiveRun>,
    pub(crate) event_tx: broadcast::Sender<EventEnvelope<RealtimeEvent>>,
    pub(crate) shutting: bool,
    /// Cancels the background lease-renewal task; owned by the actor task.
    pub(crate) renew_stop: CancellationToken,
    /// Highest realtime run_seq forwarded for the active run (batch C uses
    /// this for snapshot.last_run_seq).
    pub(crate) run_seq: Arc<AtomicU64>,
}

pub(crate) fn spawn(
    store: Arc<dyn SessionStore>,
    registry: ActorRegistry,
    session_id: SessionId,
    handle: ActorHandle,
    mut rx: mpsc::Receiver<ActorCommand>,
) {
    tokio::spawn(async move {
        let owner_id = format!("actor-{}", uuid::Uuid::new_v4());
        let owner = match store
            .acquire_lease(&session_id, &owner_id, LEASE_TTL_SECS)
            .await
        {
            Ok(owner) => owner,
            Err(error) => {
                fail_pending(rx, error, &registry, &session_id, &handle).await;
                return;
            }
        };

        match store.recover_session(&session_id, &owner).await {
            Ok(report) => {
                tracing::info!(
                    session_id = %session_id,
                    new_revision = report.new_revision,
                    interrupted = report.interrupted_runs.len(),
                    "session actor recovered session"
                );
            }
            Err(error) => {
                fail_pending(rx, error, &registry, &session_id, &handle).await;
                let _ = store.release_lease(&session_id, &owner).await;
                return;
            }
        }

        let revision = match store.get_session(&session_id).await {
            Ok(Some(session)) => session.current_revision,
            Ok(None) => {
                fail_pending(
                    rx,
                    KaynineError::SessionNotFound,
                    &registry,
                    &session_id,
                    &handle,
                )
                .await;
                let _ = store.release_lease(&session_id, &owner).await;
                return;
            }
            Err(error) => {
                fail_pending(rx, error, &registry, &session_id, &handle).await;
                let _ = store.release_lease(&session_id, &owner).await;
                return;
            }
        };

        let (event_tx, _) = broadcast::channel(BROADCAST_CAPACITY);
        let renew_stop = CancellationToken::new();
        {
            let store = store.clone();
            let session_id = session_id.clone();
            let owner = owner.clone();
            let stop = renew_stop.clone();
            tokio::spawn(async move {
                let mut ticker =
                    tokio::time::interval(std::time::Duration::from_secs(RENEW_INTERVAL_SECS));
                ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
                loop {
                    tokio::select! {
                        _ = stop.cancelled() => return,
                        _ = ticker.tick() => {}
                    }
                    match store.renew_lease(&session_id, &owner, LEASE_TTL_SECS).await {
                        Ok(true) => {}
                        Ok(false) => {
                            // A lease conflict surfaces on the next append as
                            // NotLeaseHolder/LeaseExpired; log and keep going.
                            tracing::warn!(session_id = %session_id, "writer lease renewal rejected");
                        }
                        Err(error) => {
                            tracing::warn!(session_id = %session_id, ?error, "writer lease renewal failed");
                        }
                    }
                }
            });
        }

        let mut state = ActorState {
            session_id: session_id.clone(),
            store: store.clone(),
            owner: owner.clone(),
            revision: Arc::new(tokio::sync::Mutex::new(revision)),
            active: None,
            event_tx,
            shutting: false,
            renew_stop: renew_stop.clone(),
            run_seq: Arc::new(AtomicU64::new(0)),
        };

        let mut exiting = false;
        while let Some(command) = rx.recv().await {
            match command {
                ActorCommand::StartRun { request, reply } => {
                    run::handle_start_run(&mut state, request, reply).await;
                }
                ActorCommand::Cancel { request, reply } => {
                    run::handle_cancel(&mut state, request, reply).await;
                }
                ActorCommand::Snapshot { reply } => {
                    let active = state.take_finished_active().map(|a| ActiveRunInfo {
                        run_id: a.run_id.clone(),
                        branch_id: a.branch_id.clone(),
                        model: a.model.clone(),
                    });
                    let revision = *state.revision.lock().await;
                    let run_seq = (state.run_seq.load(Ordering::SeqCst) > 0)
                        .then(|| state.run_seq.load(Ordering::SeqCst));
                    let snapshot = build_snapshot(
                        state.store.as_ref(),
                        &state.session_id,
                        active,
                        Some(revision),
                        run_seq,
                    )
                    .await;
                    let _ = reply.send(snapshot);
                }
                ActorCommand::UpdateSession { request, reply } => {
                    run::handle_update_session(&mut state, request, reply).await;
                }
                ActorCommand::Steer { request, reply } => {
                    run::handle_steer(&mut state, request, reply).await;
                }
                ActorCommand::ResolveApproval { request, reply } => {
                    run::handle_resolve_approval(&mut state, request, reply).await;
                }
                ActorCommand::Shutdown { grace, reply } => {
                    state.shutting = true;
                    let outcome = run::handle_shutdown_actor(&mut state, grace).await;
                    teardown(&mut state, &registry, &handle).await;
                    let _ = reply.send(Ok(outcome));
                    exiting = true;
                }
                ActorCommand::Release { reply } => {
                    if state.active.as_ref().is_some_and(|a| !a.join.is_finished()) {
                        let _ = reply.send(Ok(ReleaseOutcome::RunAlreadyActive));
                        continue;
                    }
                    teardown(&mut state, &registry, &handle).await;
                    let _ = reply.send(Ok(ReleaseOutcome::Released));
                    exiting = true;
                }
                ActorCommand::Subscribe { reply } => {
                    let _ = reply.send(Ok(state.event_tx.subscribe()));
                }
            }
            if exiting {
                break;
            }
        }

        teardown(&mut state, &registry, &handle).await;
    });
}

/// Stops lease renewal, releases the writer lease, and deregisters the
/// actor. Idempotent: a second call (loop-exit path after Release/Shutdown
/// already tore down) is a no-op apart from a harmless extra lease release.
async fn teardown(state: &mut ActorState, registry: &ActorRegistry, handle: &ActorHandle) {
    state.renew_stop.cancel();
    let _ = state
        .store
        .release_lease(&state.session_id, &state.owner)
        .await;
    let mut actors = registry.lock().expect("actor registry mutex poisoned");
    if actors
        .get(&state.session_id)
        .is_some_and(|current| current.actor_id == handle.actor_id)
    {
        actors.remove(&state.session_id);
    }
}

impl ActorState {
    /// Clears a finished active run (the run task appends its own terminal
    /// event, so the actor never awaits completion synchronously).
    pub(crate) fn take_finished_active(&mut self) -> Option<&ActiveRun> {
        if self.active.as_ref().is_some_and(|a| a.join.is_finished()) {
            self.active = None;
        }
        self.active.as_ref()
    }
}

/// Replies with `error` to every command already queued, then deregisters.
async fn fail_pending(
    rx: mpsc::Receiver<ActorCommand>,
    error: KaynineError,
    registry: &ActorRegistry,
    session_id: &SessionId,
    handle: &ActorHandle,
) {
    let mut rx = rx;
    while let Ok(command) = rx.try_recv() {
        match command {
            ActorCommand::StartRun { reply, .. } => {
                let _ = reply.send(Err(error.clone()));
            }
            ActorCommand::Cancel { reply, .. } => {
                let _ = reply.send(Err(error.clone()));
            }
            ActorCommand::Snapshot { reply } => {
                let _ = reply.send(Err(error.clone()));
            }
            ActorCommand::UpdateSession { reply, .. } => {
                let _ = reply.send(Err(error.clone()));
            }
            ActorCommand::Steer { reply, .. } => {
                let _ = reply.send(Err(error.clone()));
            }
            ActorCommand::ResolveApproval { reply, .. } => {
                let _ = reply.send(Err(error.clone()));
            }
            ActorCommand::Shutdown { reply, .. } => {
                let _ = reply.send(Err(error.clone()));
            }
            ActorCommand::Release { reply } => {
                let _ = reply.send(Ok(ReleaseOutcome::NotFound));
            }
            ActorCommand::Subscribe { reply } => {
                let _ = reply.send(Err(error.clone()));
            }
        }
    }
    let mut actors = registry.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);
    }
}