kcode_k1_chat_thread_durable_state/
lib.rs

1#![forbid(unsafe_code)]
2
3use kcode_k1_access_kmap::K1AccessKmap;
4use kcode_k1_chat_persistence::Session;
5pub use kcode_k1_chat_state::BoxId;
6use kcode_k1_chat_state::USER_MESSAGE_TYPE;
7use kcode_k1_chat_thread_actions::ChatThreadActions;
8pub use kcode_k1_chat_thread_actions::{AccessContext, AccessPolicy, ProfileId};
9pub use kcode_k1_chat_thread_durable_turn::{
10    BoxValue, ChatBox, PreparedCall, PreparedMailboxFlush, Status, ToolCallId,
11};
12use kcode_k1_chat_thread_durable_turn::{DurableTurn, RestartError, ShimOutput};
13use std::sync::Arc;
14
15#[derive(Clone, Debug, Eq, PartialEq)]
16pub enum TransitionError {
17    Unauthorized,
18    NotStalled,
19    NotRestartable,
20    Internal(String),
21}
22
23pub struct DurableThread {
24    turn: DurableTurn,
25    actions: ChatThreadActions,
26    authorized: bool,
27}
28
29impl DurableThread {
30    pub fn recover(session: Session, kmap: Arc<K1AccessKmap>) -> Result<Self, String> {
31        Ok(Self {
32            turn: DurableTurn::recover(session)?,
33            actions: ChatThreadActions::new(kmap),
34            authorized: false,
35        })
36    }
37
38    pub fn boxes(&self) -> &[ChatBox] {
39        self.turn.boxes()
40    }
41
42    pub fn status(&self) -> Status {
43        self.turn.status()
44    }
45
46    pub fn accept_box(
47        &mut self,
48        box_type: String,
49        contents: String,
50        hidden_type: String,
51        hidden_contents: String,
52    ) -> Result<(), String> {
53        self.turn
54            .accept(box_type, contents, hidden_type, hidden_contents)
55    }
56
57    pub fn accept_user(
58        &mut self,
59        context: AccessContext,
60        profile_id: ProfileId,
61        policy: AccessPolicy,
62        contents: String,
63    ) -> Result<(), TransitionError> {
64        let installed = self.bind_authorization(context, profile_id, policy)?;
65        match self.turn.accept(
66            USER_MESSAGE_TYPE.into(),
67            contents,
68            String::new(),
69            String::new(),
70        ) {
71            Ok(()) => Ok(()),
72            Err(error) => {
73                if installed {
74                    self.clear_authorization();
75                }
76                Err(TransitionError::Internal(error))
77            }
78        }
79    }
80
81    pub fn accept_return(
82        &mut self,
83        id: ToolCallId,
84        result: Result<String, String>,
85    ) -> Result<(), String> {
86        self.turn.accept_tool_return(id, result)
87    }
88
89    pub fn prepare_stage(
90        &mut self,
91        job: u64,
92        text: String,
93        boxes: Vec<BoxValue>,
94    ) -> Result<Vec<PreparedCall>, String> {
95        self.turn.prepare_stage(job, text, boxes)
96    }
97
98    pub fn launch_action(&mut self, name: &str, arguments: &str) -> Result<String, String> {
99        self.actions.launch(name, arguments)
100    }
101
102    pub fn accept_tool_message(&mut self, id: ToolCallId, contents: String) -> Result<(), String> {
103        self.turn.accept_tool_message(id, contents)
104    }
105
106    pub fn accept_tool_return(
107        &mut self,
108        id: ToolCallId,
109        result: Result<String, String>,
110    ) -> Result<(), String> {
111        self.turn.accept_tool_return(id, result)
112    }
113
114    pub fn accept_tool_return_v2(
115        &mut self,
116        id: ToolCallId,
117        result: Result<String, String>,
118        metadata_type: String,
119        metadata_contents: String,
120    ) -> Result<(), String> {
121        self.turn
122            .accept_tool_return_v2(id, result, metadata_type, metadata_contents)
123    }
124
125    pub fn prepare_mailbox_flush(
126        &mut self,
127        job: u64,
128    ) -> Result<Option<PreparedMailboxFlush>, String> {
129        self.turn.prepare_mailbox_flush(job)
130    }
131
132    pub fn prepared_input(&self, prepared: &PreparedMailboxFlush) -> Result<String, String> {
133        self.turn.validate_mailbox_flush(prepared)?;
134        render_input(prepared.values())
135    }
136
137    pub fn commit_mailbox_flush(&mut self, prepared: PreparedMailboxFlush) -> Result<(), String> {
138        self.turn.commit_mailbox_flush(prepared)
139    }
140
141    pub fn begin_input(&mut self) -> Result<Option<(u64, String)>, String> {
142        let Some(start) = self.turn.begin()? else {
143            return Ok(None);
144        };
145        Ok(Some((start.job, render_input(&start.values)?)))
146    }
147
148    pub fn complete(&mut self, job: u64, output: ShimOutput<BoxValue>) -> Result<bool, String> {
149        self.turn.complete(job, output)
150    }
151
152    pub fn fail(&mut self, job: u64, error: String, restartable: bool) {
153        self.turn.fail(job, error, restartable);
154        self.clear_authorization();
155    }
156
157    pub fn restart(
158        &mut self,
159        context: AccessContext,
160        profile_id: ProfileId,
161        policy: AccessPolicy,
162    ) -> Result<(), TransitionError> {
163        let installed = self.bind_authorization(context, profile_id, policy)?;
164        if let Err(error) = self.turn.restart().map_err(|error| match error {
165            RestartError::NotStalled => TransitionError::NotStalled,
166            RestartError::ProviderActionAccepted => TransitionError::NotRestartable,
167        }) {
168            if installed {
169                self.clear_authorization();
170            }
171            return Err(error);
172        }
173        Ok(())
174    }
175
176    pub fn clear_authorization(&mut self) {
177        self.actions.clear_authorization();
178        self.authorized = false;
179    }
180
181    fn bind_authorization(
182        &mut self,
183        context: AccessContext,
184        profile_id: ProfileId,
185        policy: AccessPolicy,
186    ) -> Result<bool, TransitionError> {
187        let installed = !self.authorized;
188        if self
189            .actions
190            .bind_authorization(context, profile_id, policy)
191            .is_err()
192        {
193            if installed {
194                self.actions.clear_authorization();
195            }
196            return Err(TransitionError::Unauthorized);
197        }
198        self.authorized = true;
199        Ok(installed)
200    }
201}
202
203fn render_input(values: &[BoxValue]) -> Result<String, String> {
204    let mut output = String::new();
205    for value in values {
206        let BoxValue::History(section) = value else {
207            return Err("Codex provider input contains a non-history value".into());
208        };
209        if section.is_empty() {
210            continue;
211        }
212        if !output.is_empty() && !output.ends_with('\n') {
213            output.push('\n');
214        }
215        output.push_str(section);
216    }
217    Ok(output)
218}
kcode_k1_chat_thread_durable_state/lib.rs

kcode_k1_chat_thread_durable_state/
lib.rs
