#![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_inference_settlement::settle_stage_inference;
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::AcceptStartupContext(context, profile_id, policy, items, reply) => {
Ok(reply_transition(
reply,
self.durable
.accept_startup_context(context, profile_id, policy, items),
))
}
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.is_none() || self.stage.mailbox_pending() || self.job != Some(job) {
return Err("stale Codex inference completion".to_owned());
}
drop(self.inference.take());
self.job = None;
let active_key = self.active_key.clone();
let settlement = settle_stage_inference(
&mut self.durable,
&mut self.usage,
&mut self.stage,
&active_key,
job,
result,
)
.await?;
let should_stop = settlement.should_stop();
let restore_shim = settlement.restore_shim();
let stage_error = settlement.into_stage_error();
if restore_shim {
self.shim = Some(shim);
}
if let Some(error) = stage_error {
self.stage.fail(error);
}
Ok(should_stop)
}
fn restart(
&mut self,
context: AccessContext,
profile_id: ProfileId,
policy: AccessPolicy,
) -> Result<(), ActorError> {
let generation = self
.generation
.checked_add(1)
.ok_or(ActorError::NotRestartable)?;
let active_key = format!("{}#restart-{generation}", self.base_key);
let shim = new_shim(
self.adapter.clone(),
active_key.clone(),
self.sender.clone(),
);
self.stage
.restart(&mut self.durable, context, profile_id, policy)?;
self.generation = generation;
self.active_key = active_key;
self.shim = Some(shim);
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
}