kcode_k1_chat_thread_rust_code_session/
lib.rs

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

kcode_k1_chat_thread_rust_code_session/
lib.rs
