kcode_k1_chat_thread_rust_code_session/
lib.rs1use std::sync::Arc;
2
3use kcode_k1_access_types::{AccessContext, AccessPolicy, Authority, ProfileId};
4use kcode_k1_chat_boxes::{BoxId, ChatBox, ToolCallId};
5use kcode_k1_chat_thread_durable_state::{DurableThread, PreparedMailboxFlush, TransitionError};
6use kcode_k1_chat_thread_rust_code_history::{
7 RestoredCheck as HistoryCheck, recover as recover_history,
8};
9use kcode_k1_chat_thread_rust_code_lane::{RustCodeFinished, RustCodeLane};
10use kcode_k1_chat_thread_rust_code_tasks::{RestoredCheck as TaskCheck, RustCodeTasks};
11use kcode_k1_rust_code_ktool_service::RustCodeKtoolService;
12
13pub fn is_rust_code_tool(name: &str) -> bool {
14 matches!(
15 name,
16 "RustCodeCreate"
17 | "RustCodeDocs"
18 | "RustCodeOpen"
19 | "RustCodeOverwrite"
20 | "RustCodeCheck"
21 | "RustCodePublish"
22 )
23}
24
25pub enum PersistCompletion {
26 Committed(Option<PreparedMailboxFlush>),
27 Buffered,
28}
29
30struct RetainedCompletion {
31 finished: RustCodeFinished,
32 accepted: bool,
33}
34
35pub struct RustCodeSession {
36 access: Option<AccessContext>,
37 authority: Option<Authority>,
38 lane: RustCodeLane,
39 completion: Option<RetainedCompletion>,
40 failed: bool,
41}
42
43impl RustCodeSession {
44 pub fn recover(service: Arc<RustCodeKtoolService>, boxes: &[ChatBox]) -> Result<Self, String> {
45 let tasks = match recover_history(boxes).map_err(|error| error.to_string())? {
46 None => RustCodeTasks::new(service),
47 Some(restored) => {
48 let (selector, identity, documentation, manifest, segments, ids, check) =
49 restored.into_parts();
50 let check = check.and_then(restore_check);
51 RustCodeTasks::restore(
52 service,
53 selector,
54 identity,
55 documentation,
56 manifest,
57 segments,
58 ids,
59 check,
60 )?
61 }
62 };
63 Ok(Self {
64 access: None,
65 authority: None,
66 lane: RustCodeLane::new(tasks),
67 completion: None,
68 failed: false,
69 })
70 }
71
72 pub fn accept_user(
73 &mut self,
74 durable: &mut DurableThread,
75 context: AccessContext,
76 profile_id: ProfileId,
77 policy: AccessPolicy,
78 contents: String,
79 ) -> Result<(), TransitionError> {
80 let mirrored = context.clone();
81 let authority = policy.authority();
82 durable.accept_user(context, profile_id, policy, contents)?;
83 self.access = Some(mirrored);
84 self.authority = Some(authority);
85 Ok(())
86 }
87
88 pub fn clear_authorization(&mut self, durable: &mut DurableThread) {
89 durable.clear_authorization();
90 self.clear_authorization_mirror();
91 }
92
93 pub fn clear_authorization_mirror(&mut self) {
94 self.access = None;
95 self.authority = None;
96 }
97
98 pub fn enqueue(
99 &mut self,
100 durable: &DurableThread,
101 id: ToolCallId,
102 name: String,
103 arguments: String,
104 ) -> Result<(), String> {
105 self.require_available()?;
106 if !is_rust_code_tool(&name) {
107 return Err("unsupported Rust code tool".into());
108 }
109 let access = self
110 .access
111 .clone()
112 .ok_or_else(|| "Rust code authorization is not active".to_owned())?;
113 let authority = self
114 .authority
115 .ok_or_else(|| "Rust code authorization is not active".to_owned())?;
116 let origin = locate_call(durable.boxes(), id, &name, &arguments)?;
117 self.lane
118 .enqueue(access, authority, id, origin, name, arguments)
119 .map_err(|error| error.to_string())
120 }
121
122 pub fn can_receive(&self) -> bool {
123 !self.failed && self.completion.is_none() && self.lane.is_busy()
124 }
125
126 pub fn has_completion(&self) -> bool {
127 self.completion.is_some()
128 }
129
130 pub fn is_busy(&self) -> bool {
131 !self.failed && self.lane.is_busy()
132 }
133
134 pub async fn receive(&mut self) -> Result<(), String> {
135 self.require_available()?;
136 if self.completion.is_some() {
137 return Err("a Rust code completion is already retained".into());
138 }
139 let finished = self
140 .lane
141 .receive()
142 .await
143 .map_err(|error| error.to_string())?;
144 self.completion = Some(RetainedCompletion {
145 finished,
146 accepted: false,
147 });
148 Ok(())
149 }
150
151 pub fn persist_completion(
152 &mut self,
153 durable: &mut DurableThread,
154 job: u64,
155 ) -> Result<PersistCompletion, String> {
156 self.require_available()?;
157 match self.try_persist_completion(durable, job) {
158 Ok(outcome) => Ok(outcome),
159 Err(error) => {
160 self.abort();
161 Err(error)
162 }
163 }
164 }
165
166 pub fn restart(
167 &mut self,
168 durable: &mut DurableThread,
169 service: Arc<RustCodeKtoolService>,
170 context: AccessContext,
171 profile_id: ProfileId,
172 policy: AccessPolicy,
173 ) -> Result<(), String> {
174 let mut replacement = Self::recover(service, durable.boxes())?;
175 let mirrored = context.clone();
176 let authority = policy.authority();
177 durable
178 .restart(context, profile_id, policy)
179 .map_err(|error| format!("{error:?}"))?;
180 replacement.access = Some(mirrored);
181 replacement.authority = Some(authority);
182 let mut old = std::mem::replace(self, replacement);
183 old.abort();
184 Ok(())
185 }
186
187 pub fn abort(&mut self) {
188 self.completion = None;
189 self.lane.abort();
190 self.failed = true;
191 }
192
193 fn require_available(&self) -> Result<(), String> {
194 if self.failed {
195 Err("the Rust code session is aborted".into())
196 } else {
197 Ok(())
198 }
199 }
200
201 fn try_persist_completion(
202 &mut self,
203 durable: &mut DurableThread,
204 job: u64,
205 ) -> Result<PersistCompletion, String> {
206 let (id, origin, messages, result, accepted) = {
207 let retained = self
208 .completion
209 .as_ref()
210 .ok_or_else(|| "no Rust code completion is retained".to_owned())?;
211 (
212 retained.finished.tool_call_id(),
213 retained.finished.originating_call(),
214 retained.finished.messages().to_vec(),
215 retained.finished.result().clone(),
216 retained.accepted,
217 )
218 };
219 if !accepted {
220 for message in &messages {
221 durable
222 .accept_tool_message(id, message.clone())
223 .map_err(|error| format!("failed to persist Rust Tool Message: {error}"))?;
224 }
225 durable
226 .accept_tool_return_v2(id, result.clone(), String::new(), String::new())
227 .map_err(|error| format!("failed to persist Rust Tool Result: {error}"))?;
228 self.completion
229 .as_mut()
230 .ok_or_else(|| "no Rust code completion is retained".to_owned())?
231 .accepted = true;
232 }
233 if let Some(message_ids) = correlate(durable.boxes(), id, origin, &messages, &result)? {
234 self.commit_completion(&message_ids)?;
235 return Ok(PersistCompletion::Committed(None));
236 }
237 let prepared = durable
238 .prepare_mailbox_flush(job)
239 .map_err(|error| format!("failed to prepare Rust result mailbox flush: {error}"))?;
240 let Some(message_ids) = correlate(durable.boxes(), id, origin, &messages, &result)? else {
241 return Ok(PersistCompletion::Buffered);
242 };
243 self.commit_completion(&message_ids)?;
244 Ok(PersistCompletion::Committed(prepared))
245 }
246
247 fn commit_completion(&mut self, message_ids: &[BoxId]) -> Result<(), String> {
248 let retained = self
249 .completion
250 .take()
251 .ok_or_else(|| "no Rust code completion is retained".to_owned())?;
252 self.lane
253 .commit(retained.finished, message_ids)
254 .map_err(|error| format!("failed to commit Rust code task state: {error}"))
255 }
256}
257
258#[allow(unreachable_patterns)]
259fn restore_check(value: HistoryCheck) -> Option<TaskCheck> {
260 match value {
261 HistoryCheck::Passed => Some(TaskCheck::Passed),
262 HistoryCheck::Failed(error) => Some(TaskCheck::Failed(error)),
263 _ => None,
264 }
265}
266
267fn locate_call(
268 boxes: &[ChatBox],
269 id: ToolCallId,
270 name: &str,
271 arguments: &str,
272) -> Result<BoxId, String> {
273 let mut origin = None;
274 for value in boxes {
275 let Some(call) = value
276 .tool_call_metadata()
277 .map_err(|_| "canonical Tool Call metadata is malformed".to_owned())?
278 else {
279 continue;
280 };
281 if call.tool_call_id != id {
282 continue;
283 }
284 if origin.is_some() {
285 return Err("duplicate canonical Tool Call ID".into());
286 }
287 if call.name != name || call.arguments != arguments {
288 return Err("canonical Tool Call does not match the queued call".into());
289 }
290 origin = Some(value.id());
291 }
292 origin.ok_or_else(|| "canonical Tool Call is missing".to_owned())
293}
294
295fn correlate(
296 boxes: &[ChatBox],
297 id: ToolCallId,
298 origin: BoxId,
299 messages: &[String],
300 result: &Result<String, String>,
301) -> Result<Option<Vec<BoxId>>, String> {
302 let mut found_messages = Vec::new();
303 let mut found_result = None;
304 for value in boxes {
305 if let Some(message) = value
306 .tool_message_metadata()
307 .map_err(|_| "canonical Tool Message metadata is malformed".to_owned())?
308 && message.tool_call_id == id
309 {
310 if message.originating_call != origin {
311 return Err("Rust Tool Message has the wrong origin".into());
312 }
313 found_messages.push((message.message_index, value.id(), message.message));
314 }
315 if let Some(common) = value
316 .tool_result_metadata()
317 .map_err(|_| "canonical Tool Result metadata is malformed".to_owned())?
318 && common.tool_call_id == id
319 {
320 if common.originating_call != origin {
321 return Err("Rust Tool Result has the wrong origin".into());
322 }
323 let exact = value
324 .tool_result_v2_metadata()
325 .map_err(|_| "canonical Tool Result v2 metadata is malformed".to_owned())?
326 .ok_or_else(|| "Rust completion requires Tool Result v2".to_owned())?;
327 if found_result.replace(exact).is_some() {
328 return Err("duplicate Rust Tool Result".into());
329 }
330 }
331 }
332 if found_messages.is_empty() && found_result.is_none() {
333 return Ok(None);
334 }
335 if found_messages.len() != messages.len() {
336 return Err("Rust Tool Message count does not match the completion".into());
337 }
338 found_messages.sort_by_key(|message| message.0);
339 let mut ids = Vec::with_capacity(messages.len());
340 for (offset, (index, box_id, actual)) in found_messages.into_iter().enumerate() {
341 let expected_index = u64::try_from(offset + 1)
342 .map_err(|_| "Rust Tool Message index space was exhausted".to_owned())?;
343 if index != expected_index || actual != messages[offset] {
344 return Err("Rust Tool Message correlation does not match the completion".into());
345 }
346 ids.push(box_id);
347 }
348 let actual = found_result.ok_or_else(|| "Rust Tool Result is missing".to_owned())?;
349 if actual.result != *result
350 || !actual.metadata_type.is_empty()
351 || !actual.metadata_contents.is_empty()
352 {
353 return Err("Rust Tool Result correlation does not match the completion".into());
354 }
355 Ok(Some(ids))
356}
357
358#[cfg(test)]
359mod tests {
360 use super::*;
361
362 #[test]
363 fn rust_code_name_gate_is_exact() {
364 for name in [
365 "RustCodeCreate",
366 "RustCodeDocs",
367 "RustCodeOpen",
368 "RustCodeOverwrite",
369 "RustCodeCheck",
370 "RustCodePublish",
371 ] {
372 assert!(is_rust_code_tool(name));
373 }
374 for name in ["rustcodecreate", "RustCodeWrite", "RustCodePublish ", ""] {
375 assert!(!is_rust_code_tool(name));
376 }
377 }
378
379 #[test]
380 fn known_history_checks_are_restored() {
381 assert!(matches!(
382 restore_check(HistoryCheck::Passed),
383 Some(TaskCheck::Passed)
384 ));
385 assert!(matches!(
386 restore_check(HistoryCheck::Failed("failure".into())),
387 Some(TaskCheck::Failed(error)) if error == "failure"
388 ));
389 }
390
391 #[test]
392 fn complete_package_stays_below_managed_limit() {
393 let files = [
394 include_str!("../Cargo.toml"),
395 include_str!("../Documentation.md"),
396 include_str!("lib.rs"),
397 ];
398 let count = files
399 .iter()
400 .flat_map(|file| file.lines())
401 .filter(|line| !line.trim().is_empty())
402 .count();
403 assert_eq!(count, 412, "complete package line count changed");
404 assert!(count < 500, "complete package has {count} nonblank lines");
405 }
406}