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