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