kcode-k1-chat-thread-session-provider 0.1.0

#![forbid(unsafe_code)]
#![doc = include_str!("../Documentation.md")]

use kcode_k1_chat_model_usage_session::ModelUsageSession;
use kcode_k1_chat_thread_actor_channel::{
    ActorError, ActorShim, Message, ProviderInput, ProviderInputKind, new_shim,
};
use kcode_k1_chat_thread_durable_state::{DurableThread, PreparedMailboxFlush};
use kcode_k1_chat_thread_session_view::SessionView;
use kcode_k1_codex_adapter::Adapter;
use tokio::{sync::mpsc, task::JoinHandle};

pub struct RestartPlan {
    generation: u64,
    active_key: String,
}

pub struct SessionProvider {
    adapter: Adapter,
    base_key: String,
    active_key: String,
    generation: u64,
    shim: Option<ActorShim>,
    sender: mpsc::UnboundedSender<Message>,
    inference: Option<JoinHandle<()>>,
    mailbox_flush: Option<JoinHandle<()>>,
    job: Option<u64>,
}

impl SessionProvider {
    pub fn new(
        adapter: Adapter,
        base_key: impl Into<String>,
        sender: mpsc::UnboundedSender<Message>,
    ) -> Self {
        let base_key = base_key.into();
        let active_key = base_key.clone();
        let shim = Some(new_shim(
            adapter.clone(),
            active_key.clone(),
            sender.clone(),
        ));
        Self {
            adapter,
            base_key,
            active_key,
            generation: 0,
            shim,
            sender,
            inference: None,
            mailbox_flush: None,
            job: None,
        }
    }

    pub fn job(&self) -> Option<u64> {
        self.job
    }

    pub fn active_key(&self) -> &str {
        &self.active_key
    }

    pub fn inference_active(&self) -> bool {
        self.inference.is_some()
    }

    pub fn mailbox_flush_active(&self) -> bool {
        self.mailbox_flush.is_some()
    }

    pub async fn drive(
        &mut self,
        durable: &mut DurableThread,
        usage: &mut ModelUsageSession,
        view: &mut SessionView,
    ) -> Result<(), String> {
        if self.inference.is_some() || self.shim.is_none() {
            return Ok(());
        }
        let Some((job, input)) = durable.begin_input()? else {
            return Ok(());
        };
        if !usage.is_subscribed() {
            usage
                .subscribe(&self.adapter, self.active_key.clone())
                .await
                .map_err(|_| "model-usage subscription failed".to_owned())?;
        }
        view.set_model_input(ProviderInput {
            kind: ProviderInputKind::Turn,
            text: input.clone(),
        });
        let mut shim = self.shim.take().expect("shim availability was checked");
        self.job = Some(job);
        let sender = self.sender.clone();
        self.inference = Some(tokio::spawn(async move {
            let result = shim.infer(input).await.map_err(|error| error.to_string());
            let _ = sender.send(Message::Inferred(job, shim, result));
        }));
        Ok(())
    }

    pub fn launch_mailbox_flush(
        &mut self,
        durable: &DurableThread,
        view: &mut SessionView,
        prepared: PreparedMailboxFlush,
    ) -> Result<(), String> {
        if self.mailbox_flush.is_some() {
            return Err("overlapping Codex mailbox flushes".to_owned());
        }
        let job = self
            .job
            .ok_or_else(|| "no active K1 inference".to_owned())?;
        let input = durable.prepared_input(&prepared)?;
        view.set_model_input(ProviderInput {
            kind: ProviderInputKind::MailboxFlush,
            text: input.clone(),
        });
        let adapter = self.adapter.clone();
        let key = self.active_key.clone();
        let sender = self.sender.clone();
        self.mailbox_flush = Some(tokio::spawn(async move {
            let result = adapter
                .steer(key, input)
                .await
                .map_err(|error| error.to_string());
            let _ = sender.send(Message::MailboxFlushCompleted(job, prepared, result));
        }));
        Ok(())
    }

    pub fn accept_mailbox_flush(
        &mut self,
        durable: &mut DurableThread,
        job: u64,
        prepared: PreparedMailboxFlush,
        result: Result<(), String>,
    ) -> Result<(), String> {
        if self.mailbox_flush.is_none() || self.job != Some(job) {
            return Err("stale Codex mailbox-flush completion".to_owned());
        }
        let _ = self.mailbox_flush.take();
        result?;
        durable.commit_mailbox_flush(prepared)
    }

    pub fn accept_inferred(&mut self, job: u64) -> Result<(), String> {
        if self.inference.is_none() || self.mailbox_flush.is_some() || self.job != Some(job) {
            return Err("stale Codex inference completion".to_owned());
        }
        let _ = self.inference.take();
        Ok(())
    }

    pub fn settle_success(&mut self, shim: ActorShim) {
        self.shim = Some(shim);
        self.job = None;
    }

    pub fn settle_failure(&mut self) {
        self.shim = None;
        self.job = None;
    }

    pub fn prepare_restart(&self) -> Result<RestartPlan, ActorError> {
        let generation = self
            .generation
            .checked_add(1)
            .ok_or(ActorError::NotRestartable)?;
        Ok(RestartPlan {
            generation,
            active_key: format!("{}#restart-{generation}", self.base_key),
        })
    }

    pub fn install_restart(&mut self, plan: RestartPlan) {
        self.generation = plan.generation;
        self.active_key = plan.active_key;
        self.shim = Some(new_shim(
            self.adapter.clone(),
            self.active_key.clone(),
            self.sender.clone(),
        ));
    }

    pub fn abort(&mut self) {
        for task in [self.inference.take(), self.mailbox_flush.take()]
            .into_iter()
            .flatten()
        {
            task.abort();
        }
    }
}