kcode-k1-chat-thread-durable-state 0.3.0

#![forbid(unsafe_code)]

use kcode_k1_access_kmap::K1AccessKmap;
use kcode_k1_chat_persistence::Session;
use kcode_k1_chat_state::USER_MESSAGE_TYPE;
use kcode_k1_chat_thread_actions::ChatThreadActions;
pub use kcode_k1_chat_thread_actions::{AccessContext, AccessPolicy, ProfileId};
pub use kcode_k1_chat_thread_durable_turn::{
    BoxValue, ChatBox, PreparedCall, PreparedSteer, Status, ToolCallId,
};
use kcode_k1_chat_thread_durable_turn::{DurableTurn, RestartError, ShimOutput};
use std::sync::Arc;

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum TransitionError {
    Unauthorized,
    NotStalled,
    NotRestartable,
    Internal(String),
}

pub struct DurableThread {
    turn: DurableTurn,
    actions: ChatThreadActions,
    authorized: bool,
}

impl DurableThread {
    pub fn recover(session: Session, kmap: Arc<K1AccessKmap>) -> Result<Self, String> {
        Ok(Self {
            turn: DurableTurn::recover(session)?,
            actions: ChatThreadActions::new(kmap),
            authorized: false,
        })
    }

    pub fn boxes(&self) -> &[ChatBox] {
        self.turn.boxes()
    }

    pub fn status(&self) -> Status {
        self.turn.status()
    }

    pub fn accept_box(
        &mut self,
        box_type: String,
        contents: String,
        hidden_type: String,
        hidden_contents: String,
    ) -> Result<(), String> {
        self.turn
            .accept(box_type, contents, hidden_type, hidden_contents)
    }

    pub fn accept_user(
        &mut self,
        context: AccessContext,
        profile_id: ProfileId,
        policy: AccessPolicy,
        contents: String,
    ) -> Result<(), TransitionError> {
        let installed = self.bind_authorization(context, profile_id, policy)?;
        match self.turn.accept(
            USER_MESSAGE_TYPE.into(),
            contents,
            String::new(),
            String::new(),
        ) {
            Ok(()) => Ok(()),
            Err(error) => {
                if installed {
                    self.clear_authorization();
                }
                Err(TransitionError::Internal(error))
            }
        }
    }

    pub fn accept_return(
        &mut self,
        id: ToolCallId,
        result: Result<String, String>,
    ) -> Result<(), String> {
        self.turn.accept_tool_return(id, result)
    }

    pub fn prepare_stage(
        &mut self,
        job: u64,
        text: String,
        boxes: Vec<BoxValue>,
    ) -> Result<Vec<PreparedCall>, String> {
        self.turn.prepare_stage(job, text, boxes)
    }

    pub fn launch_action(&mut self, name: &str, arguments: &str) -> Result<String, String> {
        self.actions.launch(name, arguments)
    }

    pub fn accept_tool_message(&mut self, id: ToolCallId, contents: String) -> Result<(), String> {
        self.turn.accept_tool_message(id, contents)
    }

    pub fn accept_tool_return(
        &mut self,
        id: ToolCallId,
        result: Result<String, String>,
    ) -> Result<(), String> {
        self.turn.accept_tool_return(id, result)
    }

    pub fn accept_tool_return_v2(
        &mut self,
        id: ToolCallId,
        result: Result<String, String>,
        metadata_type: String,
        metadata_contents: String,
    ) -> Result<(), String> {
        self.turn
            .accept_tool_return_v2(id, result, metadata_type, metadata_contents)
    }

    pub fn prepare_steer(&mut self, job: u64) -> Result<Option<PreparedSteer>, String> {
        self.turn.prepare_steer(job)
    }

    pub fn prepared_input(&self, prepared: &PreparedSteer) -> Result<String, String> {
        self.turn.validate_steer(prepared)?;
        render_input(prepared.values())
    }

    pub fn commit_steer(&mut self, prepared: PreparedSteer) -> Result<(), String> {
        self.turn.commit_steer(prepared)
    }

    pub fn begin_input(&mut self) -> Result<Option<(u64, String)>, String> {
        let Some(start) = self.turn.begin()? else {
            return Ok(None);
        };
        Ok(Some((start.job, render_input(&start.values)?)))
    }

    pub fn complete(&mut self, job: u64, output: ShimOutput<BoxValue>) -> Result<bool, String> {
        self.turn.complete(job, output)
    }

    pub fn fail(&mut self, job: u64, error: String, restartable: bool) {
        self.turn.fail(job, error, restartable);
        self.clear_authorization();
    }

    pub fn restart(
        &mut self,
        context: AccessContext,
        profile_id: ProfileId,
        policy: AccessPolicy,
    ) -> Result<(), TransitionError> {
        let installed = self.bind_authorization(context, profile_id, policy)?;
        if let Err(error) = self.turn.restart().map_err(|error| match error {
            RestartError::NotStalled => TransitionError::NotStalled,
            RestartError::ProviderActionAccepted => TransitionError::NotRestartable,
        }) {
            if installed {
                self.clear_authorization();
            }
            return Err(error);
        }
        Ok(())
    }

    pub fn clear_authorization(&mut self) {
        self.actions.clear_authorization();
        self.authorized = false;
    }

    fn bind_authorization(
        &mut self,
        context: AccessContext,
        profile_id: ProfileId,
        policy: AccessPolicy,
    ) -> Result<bool, TransitionError> {
        let installed = !self.authorized;
        if self
            .actions
            .bind_authorization(context, profile_id, policy)
            .is_err()
        {
            if installed {
                self.actions.clear_authorization();
            }
            return Err(TransitionError::Unauthorized);
        }
        self.authorized = true;
        Ok(installed)
    }
}

fn render_input(values: &[BoxValue]) -> Result<String, String> {
    let mut output = String::new();
    for value in values {
        let BoxValue::History(section) = value else {
            return Err("Codex provider input contains a non-history value".into());
        };
        if section.is_empty() {
            continue;
        }
        if !output.is_empty() && !output.ends_with('\n') {
            output.push('\n');
        }
        output.push_str(section);
    }
    Ok(output)
}