kcode_k1_chat_model_usage_session/

lib.rs

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

use kcode_k1_chat_model_usage::{ModelUsageCorrelation, PendingModelUsage};
use kcode_k1_chat_thread_durable_state::DurableThread;
use kcode_k1_codex_adapter::{Adapter, Error as AdapterError, RuntimeObservation};
use std::{error::Error, fmt, future};
use tokio::sync::mpsc::{self, UnboundedReceiver};

#[derive(Default)]
pub struct ModelUsageSession {
    correlation: ModelUsageCorrelation,
    key: Option<String>,
    receiver: Option<UnboundedReceiver<RuntimeObservation>>,
}

#[derive(Debug)]
pub enum ModelUsageSessionError {
    Subscription(AdapterError),
    Persistence(String),
    ConflictingKey { active: String, requested: String },
    InvalidState(&'static str),
}

impl fmt::Display for ModelUsageSessionError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Subscription(error) => write!(formatter, "usage subscription failed: {error}"),
            Self::Persistence(error) => write!(formatter, "usage persistence failed: {error}"),
            Self::ConflictingKey { active, requested } => {
                write!(
                    formatter,
                    "usage key {requested:?} conflicts with active key {active:?}"
                )
            }
            Self::InvalidState(error) => write!(formatter, "invalid usage session state: {error}"),
        }
    }
}

impl Error for ModelUsageSessionError {}

impl ModelUsageSession {
    pub async fn subscribe(
        &mut self,
        adapter: &Adapter,
        key: String,
    ) -> Result<(), ModelUsageSessionError> {
        if let Some(active) = self.active_key()? {
            return if active == key {
                Ok(())
            } else {
                Err(ModelUsageSessionError::ConflictingKey {
                    active: active.into(),
                    requested: key,
                })
            };
        }
        let receiver = adapter
            .subscribe_usage(key.clone())
            .await
            .map_err(ModelUsageSessionError::Subscription)?;
        self.key = Some(key);
        self.receiver = Some(receiver);
        Ok(())
    }

    pub async fn receive(
        &mut self,
        durable: &mut DurableThread,
    ) -> Result<bool, ModelUsageSessionError> {
        let Some(key) = self.active_key()?.map(str::to_owned) else {
            return future::pending::<Result<bool, ModelUsageSessionError>>().await;
        };
        let observation = self
            .receiver
            .as_mut()
            .ok_or(ModelUsageSessionError::InvalidState(
                "active key without receiver",
            ))?
            .recv()
            .await;
        let Some(observation) = observation else {
            self.close_receiver(durable, &key)?;
            return Ok(false);
        };
        let ready = self.correlation.observe(&key, observation);
        persist(durable, ready)?;
        Ok(true)
    }

    pub async fn drain(
        &mut self,
        durable: &mut DurableThread,
    ) -> Result<(), ModelUsageSessionError> {
        let Some(key) = self.active_key()?.map(str::to_owned) else {
            return Ok(());
        };
        loop {
            let received = self
                .receiver
                .as_mut()
                .ok_or(ModelUsageSessionError::InvalidState(
                    "active key without receiver",
                ))?
                .try_recv();
            match received {
                Ok(observation) => {
                    let ready = self.correlation.observe(&key, observation);
                    persist(durable, ready)?;
                }
                Err(mpsc::error::TryRecvError::Empty) => return Ok(()),
                Err(mpsc::error::TryRecvError::Disconnected) => {
                    self.close_receiver(durable, &key)?;
                    return Ok(());
                }
            }
        }
    }

    pub async fn finish_inference(
        &mut self,
        durable: &mut DurableThread,
        key: &str,
        terminal_agent_response_id: Option<u64>,
    ) -> Result<(), ModelUsageSessionError> {
        self.drain(durable).await?;
        self.require_active(key)?;
        let ready = self
            .correlation
            .finish_inference(key, terminal_agent_response_id);
        persist(durable, ready)
    }

    pub async fn restart(
        &mut self,
        durable: &mut DurableThread,
    ) -> Result<(), ModelUsageSessionError> {
        self.drain(durable).await?;
        if let Some(key) = self.active_key()?.map(str::to_owned) {
            let pending = self.correlation.abandon_key(&key);
            persist(durable, pending)?;
            self.clear_subscription();
        }
        Ok(())
    }

    pub async fn shutdown(
        &mut self,
        durable: &mut DurableThread,
        active_inference: bool,
    ) -> Result<(), ModelUsageSessionError> {
        self.drain(durable).await?;
        if let Some(key) = self.active_key()?.map(str::to_owned) {
            if active_inference {
                let ready = self.correlation.finish_inference(&key, None);
                persist(durable, ready)?;
            }
            let pending = self.correlation.abandon_key(&key);
            persist(durable, pending)?;
        }
        let pending = self.correlation.drain_unassociated();
        persist(durable, pending)?;
        self.clear_subscription();
        Ok(())
    }

    pub fn is_subscribed(&self) -> bool {
        self.key.is_some() && self.receiver.is_some()
    }

    fn active_key(&self) -> Result<Option<&str>, ModelUsageSessionError> {
        match (&self.key, &self.receiver) {
            (None, None) => Ok(None),
            (Some(key), Some(_)) => Ok(Some(key)),
            _ => Err(ModelUsageSessionError::InvalidState(
                "key and receiver must be installed together",
            )),
        }
    }

    fn require_active(&self, key: &str) -> Result<(), ModelUsageSessionError> {
        match self.active_key()? {
            Some(active) if active == key => Ok(()),
            Some(active) => Err(ModelUsageSessionError::ConflictingKey {
                active: active.into(),
                requested: key.into(),
            }),
            None => Err(ModelUsageSessionError::InvalidState(
                "no active usage subscription",
            )),
        }
    }

    fn close_receiver(
        &mut self,
        durable: &mut DurableThread,
        key: &str,
    ) -> Result<(), ModelUsageSessionError> {
        let pending = self.correlation.abandon_key(key);
        persist(durable, pending)?;
        self.clear_subscription();
        Ok(())
    }

    fn clear_subscription(&mut self) {
        self.key = None;
        self.receiver = None;
    }
}

fn persist(
    durable: &mut DurableThread,
    pending: Vec<PendingModelUsage>,
) -> Result<(), ModelUsageSessionError> {
    for record in pending {
        durable
            .record_model_usage(record.connected_box_id, record.usage)
            .map_err(ModelUsageSessionError::Persistence)?;
    }
    Ok(())
}

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

    #[test]
    fn default_session_is_unsubscribed() {
        assert!(!ModelUsageSession::default().is_subscribed());
    }

    #[test]
    fn conflicting_key_error_names_both_keys() {
        let error = ModelUsageSessionError::ConflictingKey {
            active: "one".into(),
            requested: "two".into(),
        };
        assert_eq!(
            error.to_string(),
            "usage key \"two\" conflicts with active key \"one\""
        );
    }
}