kcode_k1_chat_thread_durable_state/
lib.rs

1#![forbid(unsafe_code)]
2
3use kcode_k1_access_kmap::K1AccessKmap;
4use kcode_k1_chat_persistence::Session;
5pub use kcode_k1_chat_state::BoxId;
6use kcode_k1_chat_state::{AGENT_MESSAGE_TYPE, AGENT_RESPONSE_TYPE, USER_MESSAGE_TYPE};
7pub use kcode_k1_chat_thread_durable_turn::{
8    BoxValue, ChatBox, EventRecord, ModelUsage, PreflightItem, PreflightMode, PreparedCall,
9    PreparedMailboxFlush, PreparedPreflightCall, Status, TokenBreakdown, ToolCallId,
10};
11use kcode_k1_chat_thread_durable_turn::{DurableTurn, RestartError, ShimOutput};
12pub use kcode_k1_chat_thread_ktools::{AccessContext, AccessPolicy, ProfileId, SetLaunchNodeKtool};
13use kcode_k1_chat_thread_ktools::{ChatThreadKtoolExecutor, ChatThreadKtools};
14use serde::Deserialize;
15use std::sync::Arc;
16
17#[derive(Clone, Debug, Eq, PartialEq)]
18pub enum TransitionError {
19    Unauthorized,
20    NotStalled,
21    NotRestartable,
22    Internal(String),
23}
24
25pub struct DurableThread {
26    turn: DurableTurn,
27    ktools: ChatThreadKtoolExecutor,
28    authorized: bool,
29}
30
31impl DurableThread {
32    pub fn recover(session: Session, kmap: Arc<K1AccessKmap>) -> Result<Self, String> {
33        Self::recover_with_ktools(session, ChatThreadKtools::new(kmap))
34    }
35
36    pub fn recover_with_social(
37        session: Session,
38        kmap: Arc<K1AccessKmap>,
39        social: kcode_k1_ktool_social::SocialKtools,
40    ) -> Result<Self, String> {
41        Self::recover_with_ktools(session, ChatThreadKtools::new_with_social(kmap, social))
42    }
43
44    pub fn recover_with_social_and_set_launch_node(
45        session: Session,
46        kmap: Arc<K1AccessKmap>,
47        social: kcode_k1_ktool_social::SocialKtools,
48        set_launch_node: SetLaunchNodeKtool,
49    ) -> Result<Self, String> {
50        Self::recover_with_ktools(
51            session,
52            ChatThreadKtools::new_with_social_and_set_launch_node(kmap, social, set_launch_node),
53        )
54    }
55
56    pub fn boxes(&self) -> &[ChatBox] {
57        self.turn.boxes()
58    }
59
60    pub fn events(&self) -> Vec<EventRecord> {
61        self.turn.events()
62    }
63
64    pub fn status(&self) -> Status {
65        self.turn.status()
66    }
67
68    pub fn preflight_calls(&self) -> &[PreparedPreflightCall] {
69        self.turn.preflight_calls()
70    }
71
72    pub fn preflight_executor(&self) -> ChatThreadKtoolExecutor {
73        self.ktools.clone()
74    }
75
76    pub fn prepare_preflight(
77        &mut self,
78        context: AccessContext,
79        profile_id: ProfileId,
80        policy: AccessPolicy,
81        items: Vec<PreflightItem>,
82    ) -> Result<Vec<PreparedPreflightCall>, TransitionError> {
83        for item in &items {
84            if let PreflightItem::KtoolCall { name, .. } = item {
85                let supported = self
86                    .ktools
87                    .supports(name)
88                    .map_err(TransitionError::Internal)?;
89                if !supported {
90                    return Err(TransitionError::Internal(
91                        "unsupported preflight Ktool".to_owned(),
92                    ));
93                }
94            }
95        }
96        let installed = self.bind_authorization(context, profile_id, policy)?;
97        match self.turn.prepare_preflight(items) {
98            Ok(calls) => Ok(calls),
99            Err(error) => {
100                if installed {
101                    self.clear_authorization();
102                }
103                Err(TransitionError::Internal(error))
104            }
105        }
106    }
107
108    pub fn authorize_preflight(
109        &mut self,
110        context: AccessContext,
111        profile_id: ProfileId,
112        policy: AccessPolicy,
113    ) -> Result<(), TransitionError> {
114        self.bind_authorization(context, profile_id, policy)
115            .map(|_| ())
116    }
117
118    pub fn accept_box(
119        &mut self,
120        box_type: String,
121        contents: String,
122        hidden_type: String,
123        hidden_contents: String,
124    ) -> Result<(), String> {
125        self.turn
126            .accept(box_type, contents, hidden_type, hidden_contents)
127    }
128
129    pub fn accept_external_box(
130        &mut self,
131        box_type: String,
132        contents: String,
133        hidden_type: String,
134        hidden_contents: String,
135    ) -> Result<(), TransitionError> {
136        if box_type == USER_MESSAGE_TYPE {
137            return Err(TransitionError::Unauthorized);
138        }
139        self.accept_box(box_type, contents, hidden_type, hidden_contents)
140            .map_err(TransitionError::Internal)
141    }
142
143    pub fn accept_user(
144        &mut self,
145        context: AccessContext,
146        profile_id: ProfileId,
147        policy: AccessPolicy,
148        contents: String,
149    ) -> Result<(), TransitionError> {
150        let installed = self.bind_authorization(context, profile_id, policy)?;
151        match self.turn.accept(
152            USER_MESSAGE_TYPE.into(),
153            contents,
154            String::new(),
155            String::new(),
156        ) {
157            Ok(()) => Ok(()),
158            Err(error) => {
159                if installed {
160                    self.clear_authorization();
161                }
162                Err(TransitionError::Internal(error))
163            }
164        }
165    }
166
167    pub fn accept_return(
168        &mut self,
169        id: ToolCallId,
170        result: Result<String, String>,
171    ) -> Result<(), String> {
172        self.turn.accept_tool_return(id, result)
173    }
174
175    pub fn prepare_stage(
176        &mut self,
177        job: u64,
178        text: String,
179        boxes: Vec<BoxValue>,
180    ) -> Result<Vec<PreparedCall>, String> {
181        self.turn.prepare_stage(job, text, boxes)
182    }
183
184    pub fn launch_action(&mut self, name: &str, arguments: &str) -> Result<String, String> {
185        if !kcode_k1_ktool_docs::is_known_ktool(name) {
186            return Err("unknown Ktool".into());
187        }
188        match name {
189            "KtoolDocs" => kcode_k1_ktool_docs::ktool_docs(arguments),
190            "SendMessage" => launch_send_message(&mut self.turn, arguments),
191            _ => self.ktools.launch(name, arguments),
192        }
193    }
194
195    pub fn accept_tool_message(&mut self, id: ToolCallId, contents: String) -> Result<(), String> {
196        self.turn.accept_tool_message(id, contents)
197    }
198
199    pub fn accept_tool_return(
200        &mut self,
201        id: ToolCallId,
202        result: Result<String, String>,
203    ) -> Result<(), String> {
204        self.turn.accept_tool_return(id, result)
205    }
206
207    pub fn accept_tool_return_v2(
208        &mut self,
209        id: ToolCallId,
210        result: Result<String, String>,
211        metadata_type: String,
212        metadata_contents: String,
213    ) -> Result<(), String> {
214        self.turn
215            .accept_tool_return_v2(id, result, metadata_type, metadata_contents)
216    }
217
218    pub fn prepare_mailbox_flush(
219        &mut self,
220        job: u64,
221    ) -> Result<Option<PreparedMailboxFlush>, String> {
222        self.turn.prepare_mailbox_flush(job)
223    }
224
225    pub fn prepared_input(&self, prepared: &PreparedMailboxFlush) -> Result<String, String> {
226        self.turn.validate_mailbox_flush(prepared)?;
227        render_input(prepared.values())
228    }
229
230    pub fn commit_mailbox_flush(&mut self, prepared: PreparedMailboxFlush) -> Result<(), String> {
231        self.turn.commit_mailbox_flush(prepared)
232    }
233
234    pub fn begin_input(&mut self) -> Result<Option<(u64, String)>, String> {
235        let Some(start) = self.turn.begin()? else {
236            return Ok(None);
237        };
238        Ok(Some((start.job, render_input(&start.values)?)))
239    }
240
241    pub fn complete(&mut self, job: u64, output: ShimOutput<BoxValue>) -> Result<bool, String> {
242        self.turn.complete(job, output)
243    }
244
245    pub fn complete_with_terminal_response(
246        &mut self,
247        job: u64,
248        output: ShimOutput<BoxValue>,
249    ) -> Result<(bool, u64), String> {
250        let terminal_index = self.turn.boxes().len();
251        let resume = self.complete(job, output)?;
252        let terminal =
253            self.turn.boxes().get(terminal_index).ok_or_else(|| {
254                "completion did not append a terminal Agent Response box".to_owned()
255            })?;
256        if terminal.box_type() != AGENT_RESPONSE_TYPE {
257            return Err("completion terminal box was not an Agent Response".to_owned());
258        }
259        Ok((resume, terminal.id().get()))
260    }
261
262    pub fn record_model_usage(
263        &mut self,
264        connected_box_id: u64,
265        usage: ModelUsage,
266    ) -> Result<(), String> {
267        self.turn.record_model_usage(connected_box_id, usage)
268    }
269
270    pub fn fail(&mut self, job: u64, error: String, restartable: bool) {
271        self.turn.fail(job, error, restartable);
272        self.clear_authorization();
273    }
274
275    pub fn restart(
276        &mut self,
277        context: AccessContext,
278        profile_id: ProfileId,
279        policy: AccessPolicy,
280    ) -> Result<(), TransitionError> {
281        let installed = self.bind_authorization(context, profile_id, policy)?;
282        if let Err(error) = self.turn.restart().map_err(|error| match error {
283            RestartError::NotStalled => TransitionError::NotStalled,
284            RestartError::ProviderActionAccepted => TransitionError::NotRestartable,
285        }) {
286            if installed {
287                self.clear_authorization();
288            }
289            return Err(error);
290        }
291        Ok(())
292    }
293
294    pub fn clear_authorization(&mut self) {
295        if self.has_unresolved_preflight() {
296            return;
297        }
298        let _ = self.ktools.clear_authorization();
299        self.authorized = false;
300    }
301
302    fn has_unresolved_preflight(&self) -> bool {
303        self.turn.preflight_calls().iter().any(|call| {
304            !self.turn.boxes().iter().any(|value| {
305                value
306                    .tool_result_metadata()
307                    .ok()
308                    .flatten()
309                    .is_some_and(|result| result.tool_call_id == call.tool_call_id)
310            })
311        })
312    }
313
314    fn recover_with_ktools(session: Session, ktools: ChatThreadKtools) -> Result<Self, String> {
315        Ok(Self {
316            turn: DurableTurn::recover(session)?,
317            ktools: ChatThreadKtoolExecutor::new(ktools),
318            authorized: false,
319        })
320    }
321
322    fn bind_authorization(
323        &mut self,
324        context: AccessContext,
325        profile_id: ProfileId,
326        policy: AccessPolicy,
327    ) -> Result<bool, TransitionError> {
328        let installed = !self.authorized;
329        if self
330            .ktools
331            .bind_authorization(context, profile_id, policy)
332            .is_err()
333        {
334            if installed {
335                let _ = self.ktools.clear_authorization();
336            }
337            return Err(TransitionError::Unauthorized);
338        }
339        self.authorized = true;
340        Ok(installed)
341    }
342}
343
344#[derive(Deserialize)]
345#[serde(deny_unknown_fields)]
346struct SendMessageArguments {
347    message: String,
348}
349
350fn launch_send_message(turn: &mut DurableTurn, arguments: &str) -> Result<String, String> {
351    let parsed: SendMessageArguments =
352        serde_json::from_str(arguments).map_err(|_| invalid_send_message())?;
353    if parsed.message.is_empty() {
354        return Err(invalid_send_message());
355    }
356    turn.accept(
357        AGENT_MESSAGE_TYPE.into(),
358        parsed.message,
359        String::new(),
360        String::new(),
361    )?;
362    Ok("success".into())
363}
364
365fn invalid_send_message() -> String {
366    "invalid SendMessage arguments".into()
367}
368
369fn render_input(values: &[BoxValue]) -> Result<String, String> {
370    let mut output = String::new();
371    for value in values {
372        let BoxValue::History(section) = value else {
373            return Err("Codex provider input contains a non-history value".into());
374        };
375        if section.is_empty() {
376            continue;
377        }
378        if !output.is_empty() && !output.ends_with('\n') {
379            output.push('\n');
380        }
381        output.push_str(section);
382    }
383    Ok(output)
384}
385
386#[cfg(test)]
387mod tests;
kcode_k1_chat_thread_durable_state/lib.rs

kcode_k1_chat_thread_durable_state/
lib.rs
