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, PreparedCall, PreparedMailboxFlush, Status,
9    TokenBreakdown, ToolCallId,
10};
11use kcode_k1_chat_thread_durable_turn::{DurableTurn, RestartError, ShimOutput};
12use kcode_k1_chat_thread_ktools::ChatThreadKtools;
13pub use kcode_k1_chat_thread_ktools::{AccessContext, AccessPolicy, ProfileId, SetLaunchNodeKtool};
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: ChatThreadKtools,
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 accept_box(
69        &mut self,
70        box_type: String,
71        contents: String,
72        hidden_type: String,
73        hidden_contents: String,
74    ) -> Result<(), String> {
75        self.turn
76            .accept(box_type, contents, hidden_type, hidden_contents)
77    }
78
79    pub fn accept_external_box(
80        &mut self,
81        box_type: String,
82        contents: String,
83        hidden_type: String,
84        hidden_contents: String,
85    ) -> Result<(), TransitionError> {
86        if box_type == USER_MESSAGE_TYPE {
87            return Err(TransitionError::Unauthorized);
88        }
89        self.accept_box(box_type, contents, hidden_type, hidden_contents)
90            .map_err(TransitionError::Internal)
91    }
92
93    pub fn accept_user(
94        &mut self,
95        context: AccessContext,
96        profile_id: ProfileId,
97        policy: AccessPolicy,
98        contents: String,
99    ) -> Result<(), TransitionError> {
100        let installed = self.bind_authorization(context, profile_id, policy)?;
101        match self.turn.accept(
102            USER_MESSAGE_TYPE.into(),
103            contents,
104            String::new(),
105            String::new(),
106        ) {
107            Ok(()) => Ok(()),
108            Err(error) => {
109                if installed {
110                    self.clear_authorization();
111                }
112                Err(TransitionError::Internal(error))
113            }
114        }
115    }
116
117    pub fn accept_return(
118        &mut self,
119        id: ToolCallId,
120        result: Result<String, String>,
121    ) -> Result<(), String> {
122        self.turn.accept_tool_return(id, result)
123    }
124
125    pub fn prepare_stage(
126        &mut self,
127        job: u64,
128        text: String,
129        boxes: Vec<BoxValue>,
130    ) -> Result<Vec<PreparedCall>, String> {
131        self.turn.prepare_stage(job, text, boxes)
132    }
133
134    pub fn launch_action(&mut self, name: &str, arguments: &str) -> Result<String, String> {
135        if !kcode_k1_ktool_docs::is_known_ktool(name) {
136            return Err("unknown Ktool".into());
137        }
138        match name {
139            "KtoolDocs" => kcode_k1_ktool_docs::ktool_docs(arguments),
140            "SendMessage" => launch_send_message(&mut self.turn, arguments),
141            _ => self.ktools.launch(name, arguments),
142        }
143    }
144
145    pub fn accept_tool_message(&mut self, id: ToolCallId, contents: String) -> Result<(), String> {
146        self.turn.accept_tool_message(id, contents)
147    }
148
149    pub fn accept_tool_return(
150        &mut self,
151        id: ToolCallId,
152        result: Result<String, String>,
153    ) -> Result<(), String> {
154        self.turn.accept_tool_return(id, result)
155    }
156
157    pub fn accept_tool_return_v2(
158        &mut self,
159        id: ToolCallId,
160        result: Result<String, String>,
161        metadata_type: String,
162        metadata_contents: String,
163    ) -> Result<(), String> {
164        self.turn
165            .accept_tool_return_v2(id, result, metadata_type, metadata_contents)
166    }
167
168    pub fn prepare_mailbox_flush(
169        &mut self,
170        job: u64,
171    ) -> Result<Option<PreparedMailboxFlush>, String> {
172        self.turn.prepare_mailbox_flush(job)
173    }
174
175    pub fn prepared_input(&self, prepared: &PreparedMailboxFlush) -> Result<String, String> {
176        self.turn.validate_mailbox_flush(prepared)?;
177        render_input(prepared.values())
178    }
179
180    pub fn commit_mailbox_flush(&mut self, prepared: PreparedMailboxFlush) -> Result<(), String> {
181        self.turn.commit_mailbox_flush(prepared)
182    }
183
184    pub fn begin_input(&mut self) -> Result<Option<(u64, String)>, String> {
185        let Some(start) = self.turn.begin()? else {
186            return Ok(None);
187        };
188        Ok(Some((start.job, render_input(&start.values)?)))
189    }
190
191    pub fn complete(&mut self, job: u64, output: ShimOutput<BoxValue>) -> Result<bool, String> {
192        self.turn.complete(job, output)
193    }
194
195    pub fn complete_with_terminal_response(
196        &mut self,
197        job: u64,
198        output: ShimOutput<BoxValue>,
199    ) -> Result<(bool, u64), String> {
200        let terminal_index = self.turn.boxes().len();
201        let resume = self.complete(job, output)?;
202        let terminal =
203            self.turn.boxes().get(terminal_index).ok_or_else(|| {
204                "completion did not append a terminal Agent Response box".to_owned()
205            })?;
206        if terminal.box_type() != AGENT_RESPONSE_TYPE {
207            return Err("completion terminal box was not an Agent Response".to_owned());
208        }
209        Ok((resume, terminal.id().get()))
210    }
211
212    pub fn record_model_usage(
213        &mut self,
214        connected_box_id: u64,
215        usage: ModelUsage,
216    ) -> Result<(), String> {
217        self.turn.record_model_usage(connected_box_id, usage)
218    }
219
220    pub fn fail(&mut self, job: u64, error: String, restartable: bool) {
221        self.turn.fail(job, error, restartable);
222        self.clear_authorization();
223    }
224
225    pub fn restart(
226        &mut self,
227        context: AccessContext,
228        profile_id: ProfileId,
229        policy: AccessPolicy,
230    ) -> Result<(), TransitionError> {
231        let installed = self.bind_authorization(context, profile_id, policy)?;
232        if let Err(error) = self.turn.restart().map_err(|error| match error {
233            RestartError::NotStalled => TransitionError::NotStalled,
234            RestartError::ProviderActionAccepted => TransitionError::NotRestartable,
235        }) {
236            if installed {
237                self.clear_authorization();
238            }
239            return Err(error);
240        }
241        Ok(())
242    }
243
244    pub fn clear_authorization(&mut self) {
245        self.ktools.clear_authorization();
246        self.authorized = false;
247    }
248
249    fn recover_with_ktools(session: Session, ktools: ChatThreadKtools) -> Result<Self, String> {
250        Ok(Self {
251            turn: DurableTurn::recover(session)?,
252            ktools,
253            authorized: false,
254        })
255    }
256
257    fn bind_authorization(
258        &mut self,
259        context: AccessContext,
260        profile_id: ProfileId,
261        policy: AccessPolicy,
262    ) -> Result<bool, TransitionError> {
263        let installed = !self.authorized;
264        if self
265            .ktools
266            .bind_authorization(context, profile_id, policy)
267            .is_err()
268        {
269            if installed {
270                self.ktools.clear_authorization();
271            }
272            return Err(TransitionError::Unauthorized);
273        }
274        self.authorized = true;
275        Ok(installed)
276    }
277}
278
279#[derive(Deserialize)]
280#[serde(deny_unknown_fields)]
281struct SendMessageArguments {
282    message: String,
283}
284
285fn launch_send_message(turn: &mut DurableTurn, arguments: &str) -> Result<String, String> {
286    let parsed: SendMessageArguments =
287        serde_json::from_str(arguments).map_err(|_| invalid_send_message())?;
288    if parsed.message.is_empty() {
289        return Err(invalid_send_message());
290    }
291    turn.accept(
292        AGENT_MESSAGE_TYPE.into(),
293        parsed.message,
294        String::new(),
295        String::new(),
296    )?;
297    Ok("success".into())
298}
299
300fn invalid_send_message() -> String {
301    "invalid SendMessage arguments".into()
302}
303
304fn render_input(values: &[BoxValue]) -> Result<String, String> {
305    let mut output = String::new();
306    for value in values {
307        let BoxValue::History(section) = value else {
308            return Err("Codex provider input contains a non-history value".into());
309        };
310        if section.is_empty() {
311            continue;
312        }
313        if !output.is_empty() && !output.ends_with('\n') {
314            output.push('\n');
315        }
316        output.push_str(section);
317    }
318    Ok(output)
319}
320
321#[cfg(test)]
322mod tests;
kcode_k1_chat_thread_durable_state/lib.rs

kcode_k1_chat_thread_durable_state/
lib.rs
