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