syncable_ag_ui_core/
state.rs

1//! AG-UI State Management
2//!
3//! This module provides state management traits and utilities for AG-UI:
4//! - `AgentState`: Marker trait for types that can represent agent state
5//! - `FwdProps`: Marker trait for types that can be forwarded as props to UI
6//! - `StateManager`: Helper for managing state and generating deltas
7//!
8//! These traits enable generic state handling in events while ensuring
9//! the necessary bounds for serialization and async operations.
10//!
11//! # State Synchronization
12//!
13//! AG-UI supports two modes of state synchronization:
14//! - **Snapshots**: Send the complete state (simpler but less efficient)
15//! - **Deltas**: Send JSON Patch operations (more efficient for large states)
16//!
17//! The `StateManager` helper makes it easy to track state changes and
18//! generate appropriate events.
19//!
20//! # Example
21//!
22//! ```rust
23//! use ag_ui_core::state::StateManager;
24//! use serde_json::json;
25//!
26//! let mut manager = StateManager::new(json!({"count": 0}));
27//!
28//! // Update state and get the delta
29//! let delta = manager.update(json!({"count": 1}));
30//! assert!(delta.is_some());
31//!
32//! // Get current state
33//! assert_eq!(manager.current()["count"], 1);
34//! ```
35
36use crate::patch::{create_patch, Patch};
37use serde::{Deserialize, Serialize};
38use serde_json::Value as JsonValue;
39use std::fmt::Debug;
40
41/// Marker trait for types that can represent agent state.
42///
43/// Types implementing this trait can be used as the state type in
44/// state-related events (StateSnapshot, StateDelta, etc.).
45///
46/// # Bounds
47///
48/// - `'static`: Required for async operations
49/// - `Debug`: For debugging and logging
50/// - `Clone`: State may need to be copied
51/// - `Send + Sync`: For thread-safe async operations
52/// - `Serialize + Deserialize`: For JSON serialization
53/// - `Default`: For initializing empty state
54///
55/// # Example
56///
57/// ```rust
58/// use ag_ui_core::AgentState;
59/// use serde::{Deserialize, Serialize};
60///
61/// #[derive(Debug, Clone, Default, Serialize, Deserialize)]
62/// struct MyState {
63///     counter: u32,
64///     messages: Vec<String>,
65/// }
66///
67/// impl AgentState for MyState {}
68/// ```
69pub trait AgentState:
70    'static + Debug + Clone + Send + Sync + for<'de> Deserialize<'de> + Serialize + Default
71{
72}
73
74/// Marker trait for types that can be forwarded as props to UI components.
75///
76/// Types implementing this trait can be passed through the AG-UI protocol
77/// to frontend components as properties.
78///
79/// # Bounds
80///
81/// - `'static`: Required for async operations
82/// - `Clone`: Props may need to be copied
83/// - `Send + Sync`: For thread-safe async operations
84/// - `Serialize + Deserialize`: For JSON serialization
85/// - `Default`: For initializing empty props
86///
87/// # Example
88///
89/// ```rust
90/// use ag_ui_core::FwdProps;
91/// use serde::{Deserialize, Serialize};
92///
93/// #[derive(Clone, Default, Serialize, Deserialize)]
94/// struct MyProps {
95///     theme: String,
96///     locale: String,
97/// }
98///
99/// impl FwdProps for MyProps {}
100/// ```
101pub trait FwdProps:
102    'static + Clone + Send + Sync + for<'de> Deserialize<'de> + Serialize + Default
103{
104}
105
106// Implement AgentState for common types
107
108impl AgentState for JsonValue {}
109impl AgentState for () {}
110
111// Implement FwdProps for common types
112
113impl FwdProps for JsonValue {}
114impl FwdProps for () {}
115
116// =============================================================================
117// State Helper Utilities
118// =============================================================================
119
120/// Computes the difference between two JSON states as a JSON Patch.
121///
122/// Returns `None` if the states are identical.
123///
124/// # Example
125///
126/// ```rust
127/// use ag_ui_core::state::diff_states;
128/// use serde_json::json;
129///
130/// let old = json!({"count": 0});
131/// let new = json!({"count": 5});
132///
133/// let patch = diff_states(&old, &new);
134/// assert!(patch.is_some());
135/// ```
136pub fn diff_states(old: &JsonValue, new: &JsonValue) -> Option<Patch> {
137    let patch = create_patch(old, new);
138    if patch.0.is_empty() {
139        None
140    } else {
141        Some(patch)
142    }
143}
144
145/// A helper for managing state and generating deltas.
146///
147/// `StateManager` tracks the current state and provides methods to update
148/// it while automatically computing the JSON Patch delta between states.
149/// This is useful for efficiently synchronizing state with frontends.
150///
151/// # Example
152///
153/// ```rust
154/// use ag_ui_core::state::StateManager;
155/// use serde_json::json;
156///
157/// let mut manager = StateManager::new(json!({"count": 0, "items": []}));
158///
159/// // Update state - returns delta patch
160/// let delta = manager.update(json!({"count": 1, "items": []}));
161/// assert!(delta.is_some());
162///
163/// // No change - returns None
164/// let delta = manager.update(json!({"count": 1, "items": []}));
165/// assert!(delta.is_none());
166///
167/// // Check current state
168/// assert_eq!(manager.current()["count"], 1);
169/// ```
170#[derive(Debug, Clone)]
171pub struct StateManager {
172    current: JsonValue,
173    version: u64,
174}
175
176impl StateManager {
177    /// Creates a new state manager with the given initial state.
178    pub fn new(initial: JsonValue) -> Self {
179        Self {
180            current: initial,
181            version: 0,
182        }
183    }
184
185    /// Returns a reference to the current state.
186    pub fn current(&self) -> &JsonValue {
187        &self.current
188    }
189
190    /// Returns the current state version (increments on each update).
191    pub fn version(&self) -> u64 {
192        self.version
193    }
194
195    /// Updates the state and returns the delta patch if there were changes.
196    ///
197    /// Returns `None` if the new state is identical to the current state.
198    pub fn update(&mut self, new_state: JsonValue) -> Option<Patch> {
199        let patch = diff_states(&self.current, &new_state);
200        if patch.is_some() {
201            self.current = new_state;
202            self.version += 1;
203        }
204        patch
205    }
206
207    /// Updates the state using a closure and returns the delta patch.
208    ///
209    /// The closure receives a mutable reference to the current state.
210    /// After the closure completes, the delta is computed.
211    ///
212    /// # Example
213    ///
214    /// ```rust
215    /// use ag_ui_core::state::StateManager;
216    /// use serde_json::json;
217    ///
218    /// let mut manager = StateManager::new(json!({"count": 0}));
219    ///
220    /// let delta = manager.update_with(|state| {
221    ///     state["count"] = json!(10);
222    /// });
223    ///
224    /// assert!(delta.is_some());
225    /// assert_eq!(manager.current()["count"], 10);
226    /// ```
227    pub fn update_with<F>(&mut self, f: F) -> Option<Patch>
228    where
229        F: FnOnce(&mut JsonValue),
230    {
231        let old_state = self.current.clone();
232        f(&mut self.current);
233        let patch = diff_states(&old_state, &self.current);
234        if patch.is_some() {
235            self.version += 1;
236        }
237        patch
238    }
239
240    /// Resets the state to a new value without computing a delta.
241    ///
242    /// Use this when you want to replace the entire state (e.g., on reconnection)
243    /// and will send a snapshot instead of a delta.
244    pub fn reset(&mut self, new_state: JsonValue) {
245        self.current = new_state;
246        self.version += 1;
247    }
248
249    /// Takes a snapshot of the current state.
250    ///
251    /// Returns a clone of the current state value.
252    pub fn snapshot(&self) -> JsonValue {
253        self.current.clone()
254    }
255}
256
257impl Default for StateManager {
258    fn default() -> Self {
259        Self::new(JsonValue::Object(serde_json::Map::new()))
260    }
261}
262
263/// A typed state manager for custom state types.
264///
265/// This provides the same functionality as `StateManager` but works with
266/// strongly-typed state objects that implement `AgentState`.
267///
268/// # Example
269///
270/// ```rust
271/// use ag_ui_core::state::TypedStateManager;
272/// use ag_ui_core::AgentState;
273/// use serde::{Deserialize, Serialize};
274///
275/// #[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
276/// struct AppState {
277///     count: u32,
278///     user: Option<String>,
279/// }
280///
281/// impl AgentState for AppState {}
282///
283/// let mut manager = TypedStateManager::new(AppState { count: 0, user: None });
284///
285/// let delta = manager.update(AppState { count: 1, user: None });
286/// assert!(delta.is_some());
287///
288/// assert_eq!(manager.current().count, 1);
289/// ```
290#[derive(Debug, Clone)]
291pub struct TypedStateManager<S: AgentState> {
292    current: S,
293    version: u64,
294}
295
296impl<S: AgentState + PartialEq> TypedStateManager<S> {
297    /// Creates a new typed state manager with the given initial state.
298    pub fn new(initial: S) -> Self {
299        Self {
300            current: initial,
301            version: 0,
302        }
303    }
304
305    /// Returns a reference to the current state.
306    pub fn current(&self) -> &S {
307        &self.current
308    }
309
310    /// Returns the current state version (increments on each update).
311    pub fn version(&self) -> u64 {
312        self.version
313    }
314
315    /// Updates the state and returns the delta patch if there were changes.
316    ///
317    /// Returns `None` if the new state is identical to the current state.
318    pub fn update(&mut self, new_state: S) -> Option<Patch> {
319        if self.current == new_state {
320            return None;
321        }
322
323        let old_json = serde_json::to_value(&self.current).ok()?;
324        let new_json = serde_json::to_value(&new_state).ok()?;
325        let patch = diff_states(&old_json, &new_json);
326
327        self.current = new_state;
328        self.version += 1;
329        patch
330    }
331
332    /// Resets the state to a new value without computing a delta.
333    pub fn reset(&mut self, new_state: S) {
334        self.current = new_state;
335        self.version += 1;
336    }
337
338    /// Takes a snapshot of the current state as JSON.
339    pub fn snapshot(&self) -> JsonValue {
340        serde_json::to_value(&self.current).unwrap_or(JsonValue::Null)
341    }
342
343    /// Returns the current state as a JSON value.
344    pub fn as_json(&self) -> JsonValue {
345        serde_json::to_value(&self.current).unwrap_or(JsonValue::Null)
346    }
347}
348
349impl<S: AgentState + PartialEq> Default for TypedStateManager<S> {
350    fn default() -> Self {
351        Self::new(S::default())
352    }
353}
354
355#[cfg(test)]
356mod tests {
357    use super::*;
358
359    #[derive(Debug, Clone, Default, Serialize, Deserialize)]
360    struct TestState {
361        value: i32,
362    }
363
364    impl AgentState for TestState {}
365
366    #[derive(Clone, Default, Serialize, Deserialize)]
367    struct TestProps {
368        name: String,
369    }
370
371    impl FwdProps for TestProps {}
372
373    #[test]
374    fn test_json_value_implements_agent_state() {
375        fn requires_agent_state<T: AgentState>(_: T) {}
376        requires_agent_state(JsonValue::Null);
377    }
378
379    #[test]
380    fn test_unit_implements_agent_state() {
381        fn requires_agent_state<T: AgentState>(_: T) {}
382        requires_agent_state(());
383    }
384
385    #[test]
386    fn test_json_value_implements_fwd_props() {
387        fn requires_fwd_props<T: FwdProps>(_: T) {}
388        requires_fwd_props(JsonValue::Null);
389    }
390
391    #[test]
392    fn test_unit_implements_fwd_props() {
393        fn requires_fwd_props<T: FwdProps>(_: T) {}
394        requires_fwd_props(());
395    }
396
397    #[test]
398    fn test_custom_state_type() {
399        fn requires_agent_state<T: AgentState>(_: T) {}
400        requires_agent_state(TestState { value: 42 });
401    }
402
403    #[test]
404    fn test_custom_props_type() {
405        fn requires_fwd_props<T: FwdProps>(_: T) {}
406        requires_fwd_props(TestProps {
407            name: "test".to_string(),
408        });
409    }
410
411    // =========================================================================
412    // State Helper Tests
413    // =========================================================================
414
415    #[test]
416    fn test_diff_states_with_changes() {
417        use serde_json::json;
418
419        let old = json!({"count": 0});
420        let new = json!({"count": 5});
421
422        let patch = diff_states(&old, &new);
423        assert!(patch.is_some());
424    }
425
426    #[test]
427    fn test_diff_states_no_changes() {
428        use serde_json::json;
429
430        let state = json!({"count": 0});
431
432        let patch = diff_states(&state, &state);
433        assert!(patch.is_none());
434    }
435
436    #[test]
437    fn test_state_manager_new() {
438        use serde_json::json;
439
440        let manager = StateManager::new(json!({"count": 0}));
441        assert_eq!(manager.current()["count"], 0);
442        assert_eq!(manager.version(), 0);
443    }
444
445    #[test]
446    fn test_state_manager_update_with_changes() {
447        use serde_json::json;
448
449        let mut manager = StateManager::new(json!({"count": 0}));
450
451        let delta = manager.update(json!({"count": 5}));
452        assert!(delta.is_some());
453        assert_eq!(manager.current()["count"], 5);
454        assert_eq!(manager.version(), 1);
455    }
456
457    #[test]
458    fn test_state_manager_update_no_changes() {
459        use serde_json::json;
460
461        let mut manager = StateManager::new(json!({"count": 0}));
462
463        let delta = manager.update(json!({"count": 0}));
464        assert!(delta.is_none());
465        assert_eq!(manager.version(), 0); // Version shouldn't increment
466    }
467
468    #[test]
469    fn test_state_manager_update_with_closure() {
470        use serde_json::json;
471
472        let mut manager = StateManager::new(json!({"count": 0}));
473
474        let delta = manager.update_with(|state| {
475            state["count"] = json!(10);
476        });
477
478        assert!(delta.is_some());
479        assert_eq!(manager.current()["count"], 10);
480        assert_eq!(manager.version(), 1);
481    }
482
483    #[test]
484    fn test_state_manager_update_with_no_changes() {
485        use serde_json::json;
486
487        let mut manager = StateManager::new(json!({"count": 0}));
488
489        let delta = manager.update_with(|_state| {
490            // No changes
491        });
492
493        assert!(delta.is_none());
494        assert_eq!(manager.version(), 0);
495    }
496
497    #[test]
498    fn test_state_manager_reset() {
499        use serde_json::json;
500
501        let mut manager = StateManager::new(json!({"count": 0}));
502        manager.reset(json!({"count": 100, "new_field": true}));
503
504        assert_eq!(manager.current()["count"], 100);
505        assert_eq!(manager.current()["new_field"], true);
506        assert_eq!(manager.version(), 1);
507    }
508
509    #[test]
510    fn test_state_manager_snapshot() {
511        use serde_json::json;
512
513        let manager = StateManager::new(json!({"count": 42}));
514        let snapshot = manager.snapshot();
515
516        assert_eq!(snapshot, json!({"count": 42}));
517    }
518
519    #[test]
520    fn test_state_manager_default() {
521        let manager = StateManager::default();
522        assert!(manager.current().is_object());
523        assert_eq!(manager.version(), 0);
524    }
525
526    #[test]
527    fn test_state_manager_multiple_updates() {
528        use serde_json::json;
529
530        let mut manager = StateManager::new(json!({"count": 0}));
531
532        manager.update(json!({"count": 1}));
533        manager.update(json!({"count": 2}));
534        manager.update(json!({"count": 3}));
535
536        assert_eq!(manager.current()["count"], 3);
537        assert_eq!(manager.version(), 3);
538    }
539
540    // =========================================================================
541    // TypedStateManager Tests
542    // =========================================================================
543
544    #[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
545    struct AppState {
546        count: u32,
547        name: String,
548    }
549
550    impl AgentState for AppState {}
551
552    #[test]
553    fn test_typed_state_manager_new() {
554        let manager = TypedStateManager::new(AppState {
555            count: 0,
556            name: "test".to_string(),
557        });
558
559        assert_eq!(manager.current().count, 0);
560        assert_eq!(manager.current().name, "test");
561        assert_eq!(manager.version(), 0);
562    }
563
564    #[test]
565    fn test_typed_state_manager_update() {
566        let mut manager = TypedStateManager::new(AppState {
567            count: 0,
568            name: "test".to_string(),
569        });
570
571        let delta = manager.update(AppState {
572            count: 5,
573            name: "test".to_string(),
574        });
575
576        assert!(delta.is_some());
577        assert_eq!(manager.current().count, 5);
578        assert_eq!(manager.version(), 1);
579    }
580
581    #[test]
582    fn test_typed_state_manager_update_no_changes() {
583        let mut manager = TypedStateManager::new(AppState {
584            count: 0,
585            name: "test".to_string(),
586        });
587
588        let delta = manager.update(AppState {
589            count: 0,
590            name: "test".to_string(),
591        });
592
593        assert!(delta.is_none());
594        assert_eq!(manager.version(), 0);
595    }
596
597    #[test]
598    fn test_typed_state_manager_reset() {
599        let mut manager = TypedStateManager::new(AppState {
600            count: 0,
601            name: "old".to_string(),
602        });
603
604        manager.reset(AppState {
605            count: 100,
606            name: "new".to_string(),
607        });
608
609        assert_eq!(manager.current().count, 100);
610        assert_eq!(manager.current().name, "new");
611        assert_eq!(manager.version(), 1);
612    }
613
614    #[test]
615    fn test_typed_state_manager_snapshot() {
616        let manager = TypedStateManager::new(AppState {
617            count: 42,
618            name: "test".to_string(),
619        });
620
621        let snapshot = manager.snapshot();
622        assert_eq!(snapshot["count"], 42);
623        assert_eq!(snapshot["name"], "test");
624    }
625
626    #[test]
627    fn test_typed_state_manager_as_json() {
628        let manager = TypedStateManager::new(AppState {
629            count: 10,
630            name: "hello".to_string(),
631        });
632
633        let json = manager.as_json();
634        assert_eq!(json["count"], 10);
635        assert_eq!(json["name"], "hello");
636    }
637
638    #[test]
639    fn test_typed_state_manager_default() {
640        let manager: TypedStateManager<AppState> = TypedStateManager::default();
641        assert_eq!(manager.current().count, 0);
642        assert_eq!(manager.current().name, "");
643        assert_eq!(manager.version(), 0);
644    }
645}
syncable_ag_ui_core/state.rs

syncable_ag_ui_core/
state.rs
