ai_ui_core/
external_store.rs

1//! The external-store runtime, after upstream's `useExternalStoreRuntime`.
2//!
3//! The app owns the messages (and the model, transport and persistence); this
4//! runtime adapts them to the [`ThreadRuntime`] interface. The app pushes its
5//! current state with [`ExternalStoreRuntime::set_messages`] (or
6//! [`set_repository`](ExternalStoreRuntime::set_repository) for branching
7//! histories) and `set_running`; user actions arrive as [`ExternalStore`]
8//! callbacks, which the app answers by updating its messages. Nothing is
9//! inserted optimistically.
10//!
11//! Messages the app re-sends unchanged keep their `Arc` identity, so a UI can
12//! find what changed by pointer comparison without comparing content.
13
14use crate::repository::{ExportedMessage, ExportedRepository, MessageEntry, MessageRepository};
15use crate::runtime::{
16    ComposerHost, RunFuture, RuntimeCapabilities, RuntimeError, StartRunConfig, ThreadRuntime,
17    ToolResult,
18};
19use crate::subscribe::{Subscribers, Subscription};
20use crate::types::{AppendMessage, MessageId, SubmittedFeedback, ThreadMessage};
21use crate::{adapter::RunError, Composer};
22use std::cell::RefCell;
23use std::collections::HashMap;
24use std::future::Future;
25use std::pin::Pin;
26use std::rc::{Rc, Weak};
27use std::sync::Arc;
28
29/// An app callback's outcome. An error message becomes a [`RunError`].
30pub type StoreFuture = Pin<Box<dyn Future<Output = Result<(), String>>>>;
31
32/// The app side of an external-store thread (`ExternalStoreAdapter`'s
33/// `onNew`, `onEdit`, `onReload`, ...). Optional callbacks return `None` when
34/// unsupported; declare what is supported in [`capabilities`](Self::capabilities).
35pub trait ExternalStore {
36    /// A new message was sent; the app answers it by updating its messages.
37    fn on_new(&self, message: AppendMessage) -> StoreFuture;
38
39    /// A message was edited. `None` falls back to [`on_new`](Self::on_new).
40    fn on_edit(&self, _message: AppendMessage) -> Option<StoreFuture> {
41        None
42    }
43    /// Regenerate the answer to `config.parent_id`.
44    fn on_reload(&self, _config: StartRunConfig) -> Option<StoreFuture> {
45        None
46    }
47    /// Stop the run in progress.
48    fn on_cancel(&self) {}
49    /// A result was supplied for a tool call. `None` if the store does not support it.
50    fn on_add_tool_result(&self, _result: ToolResult) -> Option<StoreFuture> {
51        None
52    }
53    /// The user answered an approval request. `None` if the store does not support it.
54    fn on_respond_to_tool_approval(
55        &self,
56        _message_id: &MessageId,
57        _approval_id: &str,
58        _approved: bool,
59        _reason: Option<String>,
60    ) -> Option<StoreFuture> {
61        None
62    }
63    /// The user gave, changed or cleared feedback on a message.
64    fn on_feedback(&self, _message_id: &MessageId, _feedback: Option<&SubmittedFeedback>) {}
65    /// The user switched branch; `head_id` is the new head.
66    fn on_branch_change(&self, _head_id: Option<&MessageId>) {}
67
68    /// Which actions the UI offers. Branch switching and copy work out of the
69    /// box; the rest need the matching callbacks.
70    fn capabilities(&self) -> RuntimeCapabilities {
71        RuntimeCapabilities {
72            switch_to_branch: true,
73            edit: false,
74            reload: false,
75            cancel: false,
76            copy: true,
77            feedback: false,
78            answer_tool_call: false,
79        }
80    }
81}
82
83struct State {
84    repository: MessageRepository,
85    running: bool,
86    disabled: bool,
87}
88
89struct Inner {
90    me: Weak<Inner>,
91    store: Rc<dyn ExternalStore>,
92    state: RefCell<State>,
93    subscribers: Rc<Subscribers>,
94    thread_composer: RefCell<Option<Composer>>,
95    edits: RefCell<HashMap<MessageId, Composer>>,
96}
97
98impl Inner {
99    fn notify(&self) {
100        self.subscribers.notify();
101        let composers: Vec<Composer> = self
102            .thread_composer
103            .borrow()
104            .iter()
105            .cloned()
106            .chain(self.edits.borrow().values().cloned())
107            .collect();
108        for composer in composers {
109            composer.notify();
110        }
111    }
112}
113
114/// A thread whose messages the app owns.
115#[derive(Clone)]
116pub struct ExternalStoreRuntime {
117    inner: Rc<Inner>,
118}
119
120fn into_run(future: StoreFuture) -> RunFuture {
121    Box::pin(async move { future.await.map_err(RunError::failed) })
122}
123
124impl ExternalStoreRuntime {
125    /// A runtime over `store`, with no messages yet.
126    pub fn new(store: impl ExternalStore + 'static) -> Self {
127        Self {
128            inner: Rc::new_cyclic(|me| Inner {
129                me: me.clone(),
130                store: Rc::new(store),
131                state: RefCell::new(State {
132                    repository: MessageRepository::new(),
133                    running: false,
134                    disabled: false,
135                }),
136                subscribers: Rc::default(),
137                thread_composer: RefCell::new(None),
138                edits: RefCell::new(HashMap::new()),
139            }),
140        }
141    }
142
143    /// Called after every change. Drop the guard to unsubscribe.
144    pub fn subscribe(&self, callback: impl Fn() + 'static) -> Subscription {
145        self.inner.subscribers.subscribe(callback)
146    }
147
148    /// The app's current messages as a linear history (`messages`).
149    /// Messages equal to what is already shown keep their `Arc`.
150    pub fn set_messages<M: Into<Arc<ThreadMessage>>>(&self, messages: impl IntoIterator<Item = M>) {
151        let mut parent: Option<MessageId> = None;
152        let items = messages
153            .into_iter()
154            .map(|message| {
155                let message: Arc<ThreadMessage> = message.into();
156                let item = ExportedMessage {
157                    message: message.clone(),
158                    parent_id: parent.take(),
159                    run_config: None,
160                };
161                parent = Some(message.id().clone());
162                item
163            })
164            .collect();
165        self.set_repository(ExportedRepository {
166            head_id: None,
167            messages: items,
168        });
169    }
170
171    /// The app's messages with branches (`messageRepository`).
172    pub fn set_repository(&self, mut exported: ExportedRepository) {
173        {
174            let mut state = self.inner.state.borrow_mut();
175            // Reuse the Arc of every message that did not change.
176            for item in &mut exported.messages {
177                if let Ok(existing) = state.repository.get(item.message.id()) {
178                    if Arc::ptr_eq(&existing.message, &item.message)
179                        || *existing.message == *item.message
180                    {
181                        item.message = existing.message;
182                    }
183                }
184            }
185            // A linear history keeps the user's head only if it is still valid.
186            let keep_head = exported
187                .head_id
188                .is_none()
189                .then(|| state.repository.head_id().cloned());
190            state.repository.clear();
191            let _ = state.repository.import(exported);
192            if let Some(Some(head)) = keep_head {
193                if state.repository.contains(&head) {
194                    let _ = state.repository.switch_to_branch(&head);
195                }
196            }
197        }
198        self.inner.notify();
199    }
200
201    /// Tell the runtime whether the app is generating a response.
202    pub fn set_running(&self, running: bool) {
203        let changed = {
204            let mut state = self.inner.state.borrow_mut();
205            std::mem::replace(&mut state.running, running) != running
206        };
207        if changed {
208            self.inner.notify();
209        }
210    }
211
212    /// Disable sending (for example while the app is offline).
213    pub fn set_disabled(&self, disabled: bool) {
214        let changed = {
215            let mut state = self.inner.state.borrow_mut();
216            std::mem::replace(&mut state.disabled, disabled) != disabled
217        };
218        if changed {
219            self.inner.notify();
220        }
221    }
222
223    fn handle(inner: &Rc<Inner>) -> Self {
224        Self {
225            inner: inner.clone(),
226        }
227    }
228}
229
230impl ComposerHost for Inner {
231    fn last_message_id(&self) -> Option<MessageId> {
232        self.state
233            .borrow()
234            .repository
235            .messages()
236            .last()
237            .map(|m| m.id().clone())
238    }
239    fn append(&self, message: AppendMessage) -> Result<RunFuture, RuntimeError> {
240        let inner = self.me.upgrade().ok_or(RuntimeError::Disposed)?;
241        ExternalStoreRuntime::handle(&inner).append(message)
242    }
243    fn cancel_run(&self) {
244        self.store.on_cancel();
245    }
246    fn is_running(&self) -> bool {
247        self.state.borrow().running
248    }
249    fn end_edit(&self, id: &MessageId) {
250        self.edits.borrow_mut().remove(id);
251        self.notify();
252    }
253}
254
255impl ThreadRuntime for ExternalStoreRuntime {
256    fn messages(&self) -> Arc<[Arc<ThreadMessage>]> {
257        self.inner.state.borrow().repository.messages()
258    }
259    fn message(&self, id: &MessageId) -> Result<MessageEntry, RuntimeError> {
260        Ok(self.inner.state.borrow().repository.get(id)?)
261    }
262    fn is_running(&self) -> bool {
263        self.inner.state.borrow().running
264    }
265    fn capabilities(&self) -> RuntimeCapabilities {
266        self.inner.store.capabilities()
267    }
268    fn branches(&self, id: &MessageId) -> Result<Vec<MessageId>, RuntimeError> {
269        Ok(self.inner.state.borrow().repository.branches(id)?.to_vec())
270    }
271    fn branch_position(&self, id: &MessageId) -> Result<(usize, usize), RuntimeError> {
272        Ok(self.inner.state.borrow().repository.branch_position(id)?)
273    }
274    fn switch_to_branch(&self, id: &MessageId) -> Result<(), RuntimeError> {
275        let head = {
276            let mut state = self.inner.state.borrow_mut();
277            state.repository.switch_to_branch(id)?;
278            state.repository.head_id().cloned()
279        };
280        self.inner.notify();
281        self.inner.store.on_branch_change(head.as_ref());
282        Ok(())
283    }
284    fn append(&self, message: AppendMessage) -> Result<RunFuture, RuntimeError> {
285        if self.inner.state.borrow().disabled {
286            return Err(RuntimeError::Unsupported("sending while disabled"));
287        }
288        let editing = message.source_id.is_some();
289        let edit = editing
290            .then(|| self.inner.store.on_edit(message.clone()))
291            .flatten();
292        Ok(into_run(
293            edit.unwrap_or_else(|| self.inner.store.on_new(message)),
294        ))
295    }
296    fn reload(&self, config: StartRunConfig) -> Result<RunFuture, RuntimeError> {
297        self.inner
298            .store
299            .on_reload(config)
300            .map(into_run)
301            .ok_or(RuntimeError::Unsupported("reload"))
302    }
303    fn cancel_run(&self) {
304        self.inner.store.on_cancel();
305    }
306    fn add_tool_result(&self, result: ToolResult) -> Result<Option<RunFuture>, RuntimeError> {
307        self.inner
308            .store
309            .on_add_tool_result(result)
310            .map(|f| Some(into_run(f)))
311            .ok_or(RuntimeError::Unsupported("tool results"))
312    }
313    fn respond_to_tool_approval(
314        &self,
315        message_id: &MessageId,
316        approval_id: &str,
317        approved: bool,
318        reason: Option<String>,
319    ) -> Result<Option<RunFuture>, RuntimeError> {
320        self.inner
321            .store
322            .on_respond_to_tool_approval(message_id, approval_id, approved, reason)
323            .map(|f| Some(into_run(f)))
324            .ok_or(RuntimeError::Unsupported("tool approvals"))
325    }
326    fn submit_feedback(
327        &self,
328        id: &MessageId,
329        feedback: Option<SubmittedFeedback>,
330    ) -> Result<(), RuntimeError> {
331        if !self.inner.store.capabilities().feedback {
332            return Err(RuntimeError::Unsupported("feedback"));
333        }
334        self.message(id)?;
335        self.inner.store.on_feedback(id, feedback.as_ref());
336        Ok(())
337    }
338    fn composer(&self) -> Composer {
339        self.inner
340            .thread_composer
341            .borrow_mut()
342            .get_or_insert_with(|| {
343                Composer::thread(Rc::downgrade(&self.inner) as Weak<dyn ComposerHost>)
344            })
345            .clone()
346    }
347    fn begin_edit(&self, id: &MessageId) -> Result<Composer, RuntimeError> {
348        if self.inner.edits.borrow().contains_key(id) {
349            return Err(RuntimeError::EditInProgress(id.clone()));
350        }
351        let entry = ThreadRuntime::message(self, id)?;
352        let composer = Composer::edit(
353            Rc::downgrade(&self.inner) as Weak<dyn ComposerHost>,
354            entry.parent_id,
355            &entry.message,
356        );
357        self.inner
358            .edits
359            .borrow_mut()
360            .insert(id.clone(), composer.clone());
361        self.inner.notify();
362        Ok(composer)
363    }
364    fn edit_composer(&self, id: &MessageId) -> Option<Composer> {
365        self.inner.edits.borrow().get(id).cloned()
366    }
367    fn subscribe_dyn(&self, callback: Box<dyn Fn()>) -> Subscription {
368        self.inner.subscribers.subscribe(callback)
369    }
370    fn export(&self) -> ExportedRepository {
371        self.inner.state.borrow().repository.export()
372    }
373    fn import(&self, exported: ExportedRepository) -> Result<(), RuntimeError> {
374        self.set_repository(exported);
375        Ok(())
376    }
377}
ai_ui_core/external_store.rs

ai_ui_core/
external_store.rs
