use std::sync::Arc;
use kcode_k1_access_types::{AccessContext, AccessPolicy, ProfileId};
use kcode_k1_chat_boxes::{BoxId, ChatBox, ToolCallId};
use kcode_k1_chat_thread_durable_state::{DurableThread, PreparedMailboxFlush, TransitionError};
use kcode_k1_chat_thread_rust_code_history::{
RestoredCheck as HistoryCheck, recover as recover_history,
};
use kcode_k1_chat_thread_rust_code_lane::{RustCodeFinished, RustCodeLane};
use kcode_k1_chat_thread_rust_code_tasks::{RestoredCheck as TaskCheck, RustCodeTasks};
use kcode_k1_rust_code_ktool_service::RustCodeKtoolService;
pub fn is_rust_code_tool(name: &str) -> bool {
matches!(
name,
"RustCodeCreate"
| "RustCodeDocs"
| "RustCodeOpen"
| "RustCodeOverwrite"
| "RustCodeCheck"
| "RustCodePublish"
)
}
pub enum PersistCompletion {
Committed(Option<PreparedMailboxFlush>),
Buffered,
}
struct RetainedCompletion {
finished: RustCodeFinished,
accepted: bool,
}
pub struct RustCodeSession {
access: Option<AccessContext>,
lane: RustCodeLane,
completion: Option<RetainedCompletion>,
failed: bool,
}
impl RustCodeSession {
pub fn recover(service: Arc<RustCodeKtoolService>, boxes: &[ChatBox]) -> Result<Self, String> {
let tasks = match recover_history(boxes).map_err(|error| error.to_string())? {
None => RustCodeTasks::new(service),
Some(restored) => {
let (identity, documentation, manifest, segments, ids, check) =
restored.into_parts();
let check = check.and_then(restore_check);
RustCodeTasks::restore(
service,
identity,
documentation,
manifest,
segments,
ids,
check,
)?
}
};
Ok(Self {
access: None,
lane: RustCodeLane::new(tasks),
completion: None,
failed: false,
})
}
pub fn accept_user(
&mut self,
durable: &mut DurableThread,
context: AccessContext,
profile_id: ProfileId,
policy: AccessPolicy,
contents: String,
) -> Result<(), TransitionError> {
let mirrored = context.clone();
durable.accept_user(context, profile_id, policy, contents)?;
self.access = Some(mirrored);
Ok(())
}
pub fn clear_authorization(&mut self, durable: &mut DurableThread) {
durable.clear_authorization();
self.clear_authorization_mirror();
}
pub fn clear_authorization_mirror(&mut self) {
self.access = None;
}
pub fn enqueue(
&mut self,
durable: &DurableThread,
id: ToolCallId,
name: String,
arguments: String,
) -> Result<(), String> {
self.require_available()?;
if !is_rust_code_tool(&name) {
return Err("unsupported Rust code tool".into());
}
let access = self
.access
.clone()
.ok_or_else(|| "Rust code authorization is not active".to_owned())?;
let origin = locate_call(durable.boxes(), id, &name, &arguments)?;
self.lane
.enqueue(access, id, origin, name, arguments)
.map_err(|error| error.to_string())
}
pub fn can_receive(&self) -> bool {
!self.failed && self.completion.is_none() && self.lane.is_busy()
}
pub fn has_completion(&self) -> bool {
self.completion.is_some()
}
pub fn is_busy(&self) -> bool {
!self.failed && self.lane.is_busy()
}
pub async fn receive(&mut self) -> Result<(), String> {
self.require_available()?;
if self.completion.is_some() {
return Err("a Rust code completion is already retained".into());
}
let finished = self
.lane
.receive()
.await
.map_err(|error| error.to_string())?;
self.completion = Some(RetainedCompletion {
finished,
accepted: false,
});
Ok(())
}
pub fn persist_completion(
&mut self,
durable: &mut DurableThread,
job: u64,
) -> Result<PersistCompletion, String> {
self.require_available()?;
match self.try_persist_completion(durable, job) {
Ok(outcome) => Ok(outcome),
Err(error) => {
self.abort();
Err(error)
}
}
}
pub fn restart(
&mut self,
durable: &mut DurableThread,
service: Arc<RustCodeKtoolService>,
context: AccessContext,
profile_id: ProfileId,
policy: AccessPolicy,
) -> Result<(), String> {
let mut replacement = Self::recover(service, durable.boxes())?;
let mirrored = context.clone();
durable
.restart(context, profile_id, policy)
.map_err(|error| format!("{error:?}"))?;
replacement.access = Some(mirrored);
let mut old = std::mem::replace(self, replacement);
old.abort();
Ok(())
}
pub fn abort(&mut self) {
self.completion = None;
self.lane.abort();
self.failed = true;
}
fn require_available(&self) -> Result<(), String> {
if self.failed {
Err("the Rust code session is aborted".into())
} else {
Ok(())
}
}
fn try_persist_completion(
&mut self,
durable: &mut DurableThread,
job: u64,
) -> Result<PersistCompletion, String> {
let (id, origin, messages, result, accepted) = {
let retained = self
.completion
.as_ref()
.ok_or_else(|| "no Rust code completion is retained".to_owned())?;
(
retained.finished.tool_call_id(),
retained.finished.originating_call(),
retained.finished.messages().to_vec(),
retained.finished.result().clone(),
retained.accepted,
)
};
if !accepted {
for message in &messages {
durable
.accept_tool_message(id, message.clone())
.map_err(|error| format!("failed to persist Rust Tool Message: {error}"))?;
}
durable
.accept_tool_return_v2(id, result.clone(), String::new(), String::new())
.map_err(|error| format!("failed to persist Rust Tool Result: {error}"))?;
self.completion
.as_mut()
.ok_or_else(|| "no Rust code completion is retained".to_owned())?
.accepted = true;
}
if let Some(message_ids) = correlate(durable.boxes(), id, origin, &messages, &result)? {
self.commit_completion(&message_ids)?;
return Ok(PersistCompletion::Committed(None));
}
let prepared = durable
.prepare_mailbox_flush(job)
.map_err(|error| format!("failed to prepare Rust result mailbox flush: {error}"))?;
let Some(message_ids) = correlate(durable.boxes(), id, origin, &messages, &result)? else {
return Ok(PersistCompletion::Buffered);
};
self.commit_completion(&message_ids)?;
Ok(PersistCompletion::Committed(prepared))
}
fn commit_completion(&mut self, message_ids: &[BoxId]) -> Result<(), String> {
let retained = self
.completion
.take()
.ok_or_else(|| "no Rust code completion is retained".to_owned())?;
self.lane
.commit(retained.finished, message_ids)
.map_err(|error| format!("failed to commit Rust code task state: {error}"))
}
}
#[allow(unreachable_patterns)]
fn restore_check(value: HistoryCheck) -> Option<TaskCheck> {
match value {
HistoryCheck::Passed => Some(TaskCheck::Passed),
HistoryCheck::Failed(error) => Some(TaskCheck::Failed(error)),
_ => None,
}
}
fn locate_call(
boxes: &[ChatBox],
id: ToolCallId,
name: &str,
arguments: &str,
) -> Result<BoxId, String> {
let mut origin = None;
for value in boxes {
let Some(call) = value
.tool_call_metadata()
.map_err(|_| "canonical Tool Call metadata is malformed".to_owned())?
else {
continue;
};
if call.tool_call_id != id {
continue;
}
if origin.is_some() {
return Err("duplicate canonical Tool Call ID".into());
}
if call.name != name || call.arguments != arguments {
return Err("canonical Tool Call does not match the queued call".into());
}
origin = Some(value.id());
}
origin.ok_or_else(|| "canonical Tool Call is missing".to_owned())
}
fn correlate(
boxes: &[ChatBox],
id: ToolCallId,
origin: BoxId,
messages: &[String],
result: &Result<String, String>,
) -> Result<Option<Vec<BoxId>>, String> {
let mut found_messages = Vec::new();
let mut found_result = None;
for value in boxes {
if let Some(message) = value
.tool_message_metadata()
.map_err(|_| "canonical Tool Message metadata is malformed".to_owned())?
&& message.tool_call_id == id
{
if message.originating_call != origin {
return Err("Rust Tool Message has the wrong origin".into());
}
found_messages.push((message.message_index, value.id(), message.message));
}
if let Some(common) = value
.tool_result_metadata()
.map_err(|_| "canonical Tool Result metadata is malformed".to_owned())?
&& common.tool_call_id == id
{
if common.originating_call != origin {
return Err("Rust Tool Result has the wrong origin".into());
}
let exact = value
.tool_result_v2_metadata()
.map_err(|_| "canonical Tool Result v2 metadata is malformed".to_owned())?
.ok_or_else(|| "Rust completion requires Tool Result v2".to_owned())?;
if found_result.replace(exact).is_some() {
return Err("duplicate Rust Tool Result".into());
}
}
}
if found_messages.is_empty() && found_result.is_none() {
return Ok(None);
}
if found_messages.len() != messages.len() {
return Err("Rust Tool Message count does not match the completion".into());
}
found_messages.sort_by_key(|message| message.0);
let mut ids = Vec::with_capacity(messages.len());
for (offset, (index, box_id, actual)) in found_messages.into_iter().enumerate() {
let expected_index = u64::try_from(offset + 1)
.map_err(|_| "Rust Tool Message index space was exhausted".to_owned())?;
if index != expected_index || actual != messages[offset] {
return Err("Rust Tool Message correlation does not match the completion".into());
}
ids.push(box_id);
}
let actual = found_result.ok_or_else(|| "Rust Tool Result is missing".to_owned())?;
if actual.result != *result
|| !actual.metadata_type.is_empty()
|| !actual.metadata_contents.is_empty()
{
return Err("Rust Tool Result correlation does not match the completion".into());
}
Ok(Some(ids))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn rust_code_name_gate_is_exact() {
for name in [
"RustCodeCreate",
"RustCodeDocs",
"RustCodeOpen",
"RustCodeOverwrite",
"RustCodeCheck",
"RustCodePublish",
] {
assert!(is_rust_code_tool(name));
}
for name in ["rustcodecreate", "RustCodeWrite", "RustCodePublish ", ""] {
assert!(!is_rust_code_tool(name));
}
}
#[test]
fn known_history_checks_are_restored() {
assert!(matches!(
restore_check(HistoryCheck::Passed),
Some(TaskCheck::Passed)
));
assert!(matches!(
restore_check(HistoryCheck::Failed("failure".into())),
Some(TaskCheck::Failed(error)) if error == "failure"
));
}
#[test]
fn complete_package_stays_below_managed_limit() {
let files = [
include_str!("../Cargo.toml"),
include_str!("../Documentation.md"),
include_str!("lib.rs"),
];
let count = files
.iter()
.flat_map(|file| file.lines())
.filter(|line| !line.trim().is_empty())
.count();
assert_eq!(count, 447, "complete package line count changed");
assert!(count < 500, "complete package has {count} nonblank lines");
}
}