kcode_k1_chat_thread_session_provider/
lib.rs

1#![forbid(unsafe_code)]
2#![doc = include_str!("../Documentation.md")]
3
4use kcode_k1_chat_model_usage_session::ModelUsageSession;
5use kcode_k1_chat_thread_actor_channel::{
6    ActorError, ActorShim, Message, ProviderInput, ProviderInputKind, new_shim,
7};
8use kcode_k1_chat_thread_durable_state::{DurableThread, PreparedMailboxFlush};
9use kcode_k1_chat_thread_session_view::SessionView;
10use kcode_k1_codex_adapter::Adapter;
11use tokio::{sync::mpsc, task::JoinHandle};
12
13pub struct RestartPlan {
14    generation: u64,
15    active_key: String,
16}
17
18pub struct SessionProvider {
19    adapter: Adapter,
20    base_key: String,
21    active_key: String,
22    generation: u64,
23    shim: Option<ActorShim>,
24    sender: mpsc::UnboundedSender<Message>,
25    inference: Option<JoinHandle<()>>,
26    mailbox_flush: Option<JoinHandle<()>>,
27    job: Option<u64>,
28}
29
30impl SessionProvider {
31    pub fn new(
32        adapter: Adapter,
33        base_key: impl Into<String>,
34        sender: mpsc::UnboundedSender<Message>,
35    ) -> Self {
36        let base_key = base_key.into();
37        let active_key = base_key.clone();
38        let shim = Some(new_shim(
39            adapter.clone(),
40            active_key.clone(),
41            sender.clone(),
42        ));
43        Self {
44            adapter,
45            base_key,
46            active_key,
47            generation: 0,
48            shim,
49            sender,
50            inference: None,
51            mailbox_flush: None,
52            job: None,
53        }
54    }
55
56    pub fn job(&self) -> Option<u64> {
57        self.job
58    }
59
60    pub fn active_key(&self) -> &str {
61        &self.active_key
62    }
63
64    pub fn inference_active(&self) -> bool {
65        self.inference.is_some()
66    }
67
68    pub fn mailbox_flush_active(&self) -> bool {
69        self.mailbox_flush.is_some()
70    }
71
72    pub async fn drive(
73        &mut self,
74        durable: &mut DurableThread,
75        usage: &mut ModelUsageSession,
76        view: &mut SessionView,
77    ) -> Result<(), String> {
78        if self.inference.is_some() || self.shim.is_none() {
79            return Ok(());
80        }
81        let Some((job, input)) = durable.begin_input()? else {
82            return Ok(());
83        };
84        if !usage.is_subscribed() {
85            usage
86                .subscribe(&self.adapter, self.active_key.clone())
87                .await
88                .map_err(|_| "model-usage subscription failed".to_owned())?;
89        }
90        view.set_model_input(ProviderInput {
91            kind: ProviderInputKind::Turn,
92            text: input.clone(),
93        });
94        let mut shim = self.shim.take().expect("shim availability was checked");
95        self.job = Some(job);
96        let sender = self.sender.clone();
97        self.inference = Some(tokio::spawn(async move {
98            let result = shim.infer(input).await.map_err(|error| error.to_string());
99            let _ = sender.send(Message::Inferred(job, shim, result));
100        }));
101        Ok(())
102    }
103
104    pub fn launch_mailbox_flush(
105        &mut self,
106        durable: &DurableThread,
107        view: &mut SessionView,
108        prepared: PreparedMailboxFlush,
109    ) -> Result<(), String> {
110        if self.mailbox_flush.is_some() {
111            return Err("overlapping Codex mailbox flushes".to_owned());
112        }
113        let job = self
114            .job
115            .ok_or_else(|| "no active K1 inference".to_owned())?;
116        let input = durable.prepared_input(&prepared)?;
117        view.set_model_input(ProviderInput {
118            kind: ProviderInputKind::MailboxFlush,
119            text: input.clone(),
120        });
121        let adapter = self.adapter.clone();
122        let key = self.active_key.clone();
123        let sender = self.sender.clone();
124        self.mailbox_flush = Some(tokio::spawn(async move {
125            let result = adapter
126                .steer(key, input)
127                .await
128                .map_err(|error| error.to_string());
129            let _ = sender.send(Message::MailboxFlushCompleted(job, prepared, result));
130        }));
131        Ok(())
132    }
133
134    pub fn accept_mailbox_flush(
135        &mut self,
136        durable: &mut DurableThread,
137        job: u64,
138        prepared: PreparedMailboxFlush,
139        result: Result<(), String>,
140    ) -> Result<(), String> {
141        if self.mailbox_flush.is_none() || self.job != Some(job) {
142            return Err("stale Codex mailbox-flush completion".to_owned());
143        }
144        let _ = self.mailbox_flush.take();
145        result?;
146        durable.commit_mailbox_flush(prepared)
147    }
148
149    pub fn accept_inferred(&mut self, job: u64) -> Result<(), String> {
150        if self.inference.is_none() || self.mailbox_flush.is_some() || self.job != Some(job) {
151            return Err("stale Codex inference completion".to_owned());
152        }
153        let _ = self.inference.take();
154        Ok(())
155    }
156
157    pub fn settle_success(&mut self, shim: ActorShim) {
158        self.shim = Some(shim);
159        self.job = None;
160    }
161
162    pub fn settle_failure(&mut self) {
163        self.shim = None;
164        self.job = None;
165    }
166
167    pub fn prepare_restart(&self) -> Result<RestartPlan, ActorError> {
168        let generation = self
169            .generation
170            .checked_add(1)
171            .ok_or(ActorError::NotRestartable)?;
172        Ok(RestartPlan {
173            generation,
174            active_key: format!("{}#restart-{generation}", self.base_key),
175        })
176    }
177
178    pub fn install_restart(&mut self, plan: RestartPlan) {
179        self.generation = plan.generation;
180        self.active_key = plan.active_key;
181        self.shim = Some(new_shim(
182            self.adapter.clone(),
183            self.active_key.clone(),
184            self.sender.clone(),
185        ));
186    }
187
188    pub fn abort(&mut self) {
189        for task in [self.inference.take(), self.mailbox_flush.take()]
190            .into_iter()
191            .flatten()
192        {
193            task.abort();
194        }
195    }
196}
kcode_k1_chat_thread_session_provider/lib.rs

kcode_k1_chat_thread_session_provider/
lib.rs
