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