kcode_k1_chat_thread_session_actor/

lib.rs

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

use kcode_k1_access_kmap::K1AccessKmap;
use kcode_k1_chat_model_usage_session::ModelUsageSession;
use kcode_k1_chat_persistence::Session;
use kcode_k1_chat_thread_actor_channel::{
    ActorError, ActorShim, BoxValue, Handle, Message, ProviderInput, ProviderInputKind, Reply,
    channel_with_events, new_shim,
};
use kcode_k1_chat_thread_durable_state::{
    AccessContext, AccessPolicy, DurableThread, ProfileId, SetLaunchNodeKtool, TransitionError,
};
use kcode_k1_chat_thread_session_code_runtime::SessionCodeRuntime;
use kcode_k1_chat_thread_session_stage_runtime::{EventContext, StageRuntime};
use kcode_k1_chat_thread_session_view::SessionView;
use kcode_k1_codex_adapter::{Adapter, ShimOutput};
use kcode_k1_codex_websearch::Runner as WebSearchRunner;
use kcode_k1_rust_code_ktool_service::RustCodeKtoolService;
use kcode_k1_web_code_ktool_service::K1WebCodeKtoolService;
use std::sync::Arc;
use tokio::{sync::mpsc, task::JoinHandle};

type ActorResult = Result<bool, String>;
type InferenceResult = Result<ShimOutput<BoxValue>, String>;
type UnitResult = Result<(), String>;

pub fn open(
    adapter: Adapter,
    key: impl Into<String>,
    session: Session,
    kmap: Arc<K1AccessKmap>,
    web_search: WebSearchRunner,
) -> Result<Handle, String> {
    let durable = DurableThread::recover(session, kmap)?;
    spawn_actor(adapter, key, durable, None, None, web_search)
}

pub fn open_with_social(
    adapter: Adapter,
    key: impl Into<String>,
    session: Session,
    kmap: Arc<K1AccessKmap>,
    social: kcode_k1_ktool_social::SocialKtools,
    web_search: WebSearchRunner,
) -> Result<Handle, String> {
    let durable = DurableThread::recover_with_social(session, kmap, social)?;
    spawn_actor(adapter, key, durable, None, None, web_search)
}

pub fn open_with_social_and_set_launch_node(
    adapter: Adapter,
    key: impl Into<String>,
    session: Session,
    kmap: Arc<K1AccessKmap>,
    social: kcode_k1_ktool_social::SocialKtools,
    set_launch_node: SetLaunchNodeKtool,
    web_search: WebSearchRunner,
) -> Result<Handle, String> {
    let durable = DurableThread::recover_with_social_and_set_launch_node(
        session,
        kmap,
        social,
        set_launch_node,
    )?;
    spawn_actor(adapter, key, durable, None, None, web_search)
}

#[allow(clippy::too_many_arguments)]
pub fn open_with_social_and_set_launch_node_and_rust_code(
    adapter: Adapter,
    key: impl Into<String>,
    session: Session,
    kmap: Arc<K1AccessKmap>,
    social: kcode_k1_ktool_social::SocialKtools,
    set_launch_node: SetLaunchNodeKtool,
    rust_code: Arc<RustCodeKtoolService>,
    web_search: WebSearchRunner,
) -> Result<Handle, String> {
    let durable = DurableThread::recover_with_social_and_set_launch_node(
        session,
        kmap,
        social,
        set_launch_node,
    )?;
    spawn_actor(adapter, key, durable, Some(rust_code), None, web_search)
}

#[allow(clippy::too_many_arguments)]
pub fn open_with_social_and_set_launch_node_and_rust_code_and_web_code(
    adapter: Adapter,
    key: impl Into<String>,
    session: Session,
    kmap: Arc<K1AccessKmap>,
    social: kcode_k1_ktool_social::SocialKtools,
    set_launch_node: SetLaunchNodeKtool,
    rust_code: Arc<RustCodeKtoolService>,
    web_code: K1WebCodeKtoolService,
    web_search: WebSearchRunner,
) -> Result<Handle, String> {
    let durable = DurableThread::recover_with_social_and_set_launch_node(
        session,
        kmap,
        social,
        set_launch_node,
    )?;
    spawn_actor(
        adapter,
        key,
        durable,
        Some(rust_code),
        Some(web_code),
        web_search,
    )
}

fn spawn_actor(
    adapter: Adapter,
    key: impl Into<String>,
    durable: DurableThread,
    rust_code: Option<Arc<RustCodeKtoolService>>,
    web_code: Option<K1WebCodeKtoolService>,
    web_search: WebSearchRunner,
) -> Result<Handle, String> {
    let code = SessionCodeRuntime::recover(&durable, rust_code, web_code)?;
    let stage = StageRuntime::new(code, web_search);
    let (handle, sender, receiver, event_receiver) = channel_with_events();
    let base_key = key.into();
    let actor = Actor {
        durable,
        shim: Some(new_shim(adapter.clone(), base_key.clone(), sender.clone())),
        adapter,
        active_key: base_key.clone(),
        base_key,
        generation: 0,
        stage,
        view: SessionView::new(event_receiver),
        sender,
        receiver,
        usage: ModelUsageSession::default(),
        inference: None,
        job: None,
    };
    tokio::spawn(actor.run());
    Ok(handle)
}

struct Actor {
    durable: DurableThread,
    shim: Option<ActorShim>,
    adapter: Adapter,
    base_key: String,
    active_key: String,
    generation: u64,
    stage: StageRuntime,
    view: SessionView,
    sender: mpsc::UnboundedSender<Message>,
    receiver: mpsc::UnboundedReceiver<Message>,
    usage: ModelUsageSession,
    inference: Option<JoinHandle<()>>,
    job: Option<u64>,
}

impl Actor {
    async fn run(mut self) {
        loop {
            if let Err(error) = self.drive().await {
                self.fatal(error);
                break;
            }
            let searching = self.stage.search_active();
            let can_receive_code = self.stage.can_receive_code();
            self.view.wake(&self.durable, searching);
            let stop = tokio::select! {
                message = self.receiver.recv() => match message {
                    Some(message) => match self.handle(message).await {
                        Ok(stop) => stop,
                        Err(error) => self.fatal(error),
                    },
                    None => true,
                },
                result = self.stage.receive_code(), if can_receive_code => match result {
                    Ok(()) => match self.handle_code_completion() {
                        Ok(stop) => stop,
                        Err(error) => self.fatal(error),
                    },
                    Err(error) => self.fatal(error),
                },
                reply = self.view.receive_event_query() => {
                    if let Some(reply) = reply {
                        self.view.answer_event_query(&self.durable, searching, reply);
                    }
                    false
                }
                usage = self.usage.receive(&mut self.durable) => usage.is_err(),
            };
            if stop {
                break;
            }
        }
        self.stage.abort();
        let inference_active = self.inference.is_some();
        let _ = self
            .usage
            .shutdown(&mut self.durable, inference_active)
            .await;
        self.stage.shutdown();
        if let Some(task) = self.inference.take() {
            task.abort();
        }
        self.view.close();
    }

    async fn handle(&mut self, message: Message) -> ActorResult {
        match message {
            Message::Accept((box_type, contents, hidden_type, hidden_contents), reply) => {
                Ok(reply_transition(
                    reply,
                    self.durable.accept_external_box(
                        box_type,
                        contents,
                        hidden_type,
                        hidden_contents,
                    ),
                ))
            }
            Message::AcceptUser(context, profile_id, policy, contents, reply) => {
                let result = self.stage.accept_user(
                    &mut self.durable,
                    context,
                    profile_id,
                    policy,
                    contents,
                );
                Ok(reply_transition(reply, result))
            }
            Message::Return(id, result, reply) => {
                let result = self
                    .durable
                    .accept_return(id, result)
                    .map_err(|_| ActorError::Closed);
                let stop = result.is_err();
                Ok(answer(reply, result, stop))
            }
            Message::Stage(text, boxes, reply) => {
                let mut event = EventContext {
                    durable: &mut self.durable,
                    adapter: &self.adapter,
                    active_key: &self.active_key,
                    view: &mut self.view,
                    sender: &self.sender,
                    job: self.job,
                };
                self.stage.handle_stage(&mut event, text, boxes, reply)
            }
            Message::WebSearchCompleted(epoch, id, result) => {
                let mut event = EventContext {
                    durable: &mut self.durable,
                    adapter: &self.adapter,
                    active_key: &self.active_key,
                    view: &mut self.view,
                    sender: &self.sender,
                    job: self.job,
                };
                self.stage
                    .handle_web_search_completion(&mut event, epoch, id, result)
            }
            Message::MailboxFlushCompleted(job, prepared, result) => {
                let mut event = EventContext {
                    durable: &mut self.durable,
                    adapter: &self.adapter,
                    active_key: &self.active_key,
                    view: &mut self.view,
                    sender: &self.sender,
                    job: self.job,
                };
                self.stage
                    .handle_mailbox_flush_completion(&mut event, job, prepared, result)
            }
            Message::Inferred(job, shim, result) => self.finish(job, shim, result).await,
            Message::Snapshot(reply) => {
                let snapshot = self
                    .view
                    .snapshot(&self.durable, self.stage.search_active());
                Ok(answer(reply, Ok(snapshot), false))
            }
            Message::Wait(reply) => {
                self.view
                    .wait(&self.durable, self.stage.search_active(), reply);
                Ok(false)
            }
            Message::Restart(context, profile_id, policy, reply) => {
                let result = self.restart(context, profile_id, policy);
                if let Err(error) = result {
                    return Ok(answer(reply, Err(error), false));
                }
                if self.usage.restart(&mut self.durable).await.is_err() {
                    Ok(answer(reply, Err(ActorError::Closed), true))
                } else {
                    Ok(answer(reply, Ok(()), false))
                }
            }
            Message::Abandon => Ok(true),
        }
    }

    fn handle_code_completion(&mut self) -> ActorResult {
        let mut event = EventContext {
            durable: &mut self.durable,
            adapter: &self.adapter,
            active_key: &self.active_key,
            view: &mut self.view,
            sender: &self.sender,
            job: self.job,
        };
        self.stage.handle_code_completion(&mut event)
    }

    fn fatal(&mut self, error: String) -> bool {
        self.stage.fail(error)
    }

    async fn drive(&mut self) -> UnitResult {
        if self.inference.is_some() || self.shim.is_none() {
            return Ok(());
        }
        let Some((job, input)) = self.durable.begin_input()? else {
            return Ok(());
        };
        if !self.usage.is_subscribed() {
            self.usage
                .subscribe(&self.adapter, self.active_key.clone())
                .await
                .map_err(|_| "model-usage subscription failed".to_owned())?;
        }
        self.view.set_model_input(ProviderInput {
            kind: ProviderInputKind::Turn,
            text: input.clone(),
        });
        let mut shim = self.shim.take().expect("shim 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(())
    }

    async fn finish(&mut self, job: u64, shim: ActorShim, result: InferenceResult) -> ActorResult {
        if self.inference.take().is_none() || self.stage.mailbox_pending() || self.job != Some(job)
        {
            return Err("stale Codex inference completion".to_owned());
        }
        if self.stage.code_pending() {
            return Err("Codex inference completed while Rust code work remains".to_owned());
        }
        self.job = None;
        let key = self.active_key.clone();
        let (error, resume, terminal_id) = match result {
            Err(error) => {
                self.stage.abort();
                (Some(error), false, None)
            }
            Ok(output) => {
                let (resume, terminal) =
                    self.durable.complete_with_terminal_response(job, output)?;
                (None, resume, Some(terminal))
            }
        };
        if self
            .usage
            .finish_inference(&mut self.durable, &key, terminal_id)
            .await
            .is_err()
        {
            return Ok(true);
        }
        if let Some(error) = error {
            self.stage.fail(error.clone());
            self.durable.fail(job, error, true);
            return Ok(false);
        }
        self.shim = Some(shim);
        if !resume && !self.stage.search_active() {
            self.stage.clear_authorization(&mut self.durable);
        }
        Ok(false)
    }

    fn restart(
        &mut self,
        context: AccessContext,
        profile_id: ProfileId,
        policy: AccessPolicy,
    ) -> Result<(), ActorError> {
        let generation = self
            .generation
            .checked_add(1)
            .ok_or(ActorError::NotRestartable)?;
        self.stage
            .restart(&mut self.durable, context, profile_id, policy)?;
        self.generation = generation;
        self.active_key = format!("{}#restart-{generation}", self.base_key);
        self.shim = Some(new_shim(
            self.adapter.clone(),
            self.active_key.clone(),
            self.sender.clone(),
        ));
        Ok(())
    }
}

fn map_transition(error: TransitionError) -> ActorError {
    match error {
        TransitionError::Unauthorized => ActorError::Unauthorized,
        TransitionError::NotStalled => ActorError::NotStalled,
        TransitionError::NotRestartable => ActorError::NotRestartable,
        TransitionError::Internal(_) => ActorError::Closed,
    }
}

fn reply_transition(reply: Reply<()>, result: Result<(), TransitionError>) -> bool {
    let stop = matches!(result, Err(TransitionError::Internal(_)));
    answer(reply, result.map_err(map_transition), stop)
}

fn answer<T>(reply: Reply<T>, result: Result<T, ActorError>, stop: bool) -> bool {
    let _ = reply.send(result);
    stop
}