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