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

kcode_k1_chat_thread_rust_code_session/
lib.rs
