1+ use std:: error:: Error ;
2+ use std:: fmt:: { Debug , Display , Formatter } ;
3+ use std:: sync:: Arc ;
4+
5+ use async_trait:: async_trait;
6+ use log:: info;
7+ use reqwest:: header:: HeaderMap ;
8+ use reqwest:: Url ;
9+ use serde:: { Deserialize , Serialize } ;
10+
11+ use ag_ui_core:: { AgentState , FwdProps , JsonValue } ;
12+ use ag_ui_core:: event:: { StateDeltaEvent , StateSnapshotEvent } ;
13+ use ag_ui_core:: types:: ids:: MessageId ;
14+ use ag_ui_core:: types:: message:: Message ;
15+ use ag_ui_client:: agent:: { AgentError , AgentStateMutation , RunAgentParams } ;
16+ use ag_ui_client:: { Agent , HttpAgent } ;
17+ use ag_ui_client:: subscriber:: { AgentSubscriber , AgentSubscriberParams } ;
18+
19+ #[ derive( Serialize , Deserialize , Debug , Clone , PartialEq , Eq , Hash ) ]
20+ pub enum StepStatus {
21+ #[ serde( rename = "pending" ) ]
22+ Pending ,
23+ #[ serde( rename = "completed" ) ]
24+ Completed ,
25+ }
26+
27+ impl Display for StepStatus {
28+ fn fmt ( & self , f : & mut Formatter < ' _ > ) -> std:: fmt:: Result {
29+ match self {
30+ StepStatus :: Pending => write ! ( f, "pending" ) ,
31+ StepStatus :: Completed => write ! ( f, "completed" ) ,
32+ }
33+ }
34+ }
35+
36+ impl Default for StepStatus {
37+ fn default ( ) -> Self {
38+ StepStatus :: Pending
39+ }
40+ }
41+
42+ #[ derive( Serialize , Deserialize , Debug , Clone ) ]
43+ pub struct Step {
44+ pub description : String ,
45+ #[ serde( default ) ]
46+ pub status : StepStatus ,
47+ }
48+
49+ impl Step {
50+ pub fn new ( description : String ) -> Self {
51+ Self {
52+ description,
53+ status : StepStatus :: Pending ,
54+ }
55+ }
56+ }
57+
58+ #[ derive( Serialize , Deserialize , Debug , Clone , Default ) ]
59+ pub struct Plan {
60+ #[ serde( default ) ]
61+ pub steps : Vec < Step > ,
62+ }
63+
64+ impl AgentState for Plan { }
65+
66+ pub struct GenerativeUiSubscriber ;
67+
68+ impl GenerativeUiSubscriber {
69+ pub fn new ( ) -> Self {
70+ Self
71+ }
72+ }
73+
74+ #[ async_trait]
75+ impl < FwdPropsT > AgentSubscriber < Plan , FwdPropsT > for GenerativeUiSubscriber
76+ where
77+ FwdPropsT : FwdProps + Debug ,
78+ {
79+ async fn on_state_snapshot_event (
80+ & self ,
81+ event : & StateSnapshotEvent < Plan > ,
82+ _params : AgentSubscriberParams < ' async_trait , Plan , FwdPropsT > ,
83+ ) -> Result < AgentStateMutation < Plan > , AgentError > {
84+ info ! ( "📸 State snapshot received:" ) ;
85+ let plan = & event. snapshot ;
86+ info ! ( " Plan with {} steps:" , plan. steps. len( ) ) ;
87+ for ( i, step) in plan. steps . iter ( ) . enumerate ( ) {
88+ let status_icon = match step. status {
89+ StepStatus :: Pending => "⏳" ,
90+ StepStatus :: Completed => "✅" ,
91+ } ;
92+ info ! ( " {}. {} {}" , i + 1 , status_icon, step. description) ;
93+ }
94+ Ok ( AgentStateMutation :: default ( ) )
95+ }
96+
97+ async fn on_state_delta_event (
98+ & self ,
99+ event : & StateDeltaEvent ,
100+ _params : AgentSubscriberParams < ' async_trait , Plan , FwdPropsT >
101+ ) -> Result < AgentStateMutation < Plan > , AgentError > {
102+ info ! ( "🔄 State delta received:" ) ;
103+ for patch in & event. delta {
104+ match patch. get ( "op" ) . and_then ( |v| v. as_str ( ) ) {
105+ Some ( "replace" ) => {
106+ if let ( Some ( path) , Some ( value) ) = (
107+ patch. get ( "path" ) . and_then ( |v| v. as_str ( ) ) ,
108+ patch. get ( "value" )
109+ ) {
110+ if path. contains ( "/status" ) {
111+ let status = value. as_str ( ) . unwrap_or ( "unknown" ) ;
112+ let status_icon = match status {
113+ "completed" => "✅" ,
114+ "pending" => "⏳" ,
115+ _ => "❓" ,
116+ } ;
117+ info ! ( " {} Step status updated to: {}" , status_icon, status) ;
118+ } else if path. contains ( "/description" ) {
119+ info ! ( " 📝 Step description updated to: {}" , value. as_str( ) . unwrap_or( "unknown" ) ) ;
120+ }
121+ }
122+ }
123+ Some ( op) => info ! ( " Operation: {}" , op) ,
124+ None => info ! ( " Unknown operation" ) ,
125+ }
126+ }
127+ Ok ( AgentStateMutation :: default ( ) )
128+ }
129+
130+ async fn on_state_changed (
131+ & self ,
132+ params : AgentSubscriberParams < ' async_trait , Plan , FwdPropsT >
133+ ) -> Result < ( ) , AgentError > {
134+ info ! ( "🔄 Overall state changed" ) ;
135+ let completed_steps = params. state . steps . iter ( )
136+ . filter ( |step| matches ! ( step. status, StepStatus :: Completed ) )
137+ . count ( ) ;
138+ info ! ( " Progress: {}/{} steps completed" , completed_steps, params. state. steps. len( ) ) ;
139+
140+ Ok ( ( ) )
141+ }
142+ }
143+
144+ #[ tokio:: main]
145+ async fn main ( ) -> Result < ( ) , Box < dyn Error > > {
146+ env_logger:: Builder :: from_env ( env_logger:: Env :: default ( ) . default_filter_or ( "info" ) ) . init ( ) ;
147+
148+ let base_url = Url :: parse ( "http://127.0.0.1:3001/" ) ?;
149+ let headers = HeaderMap :: new ( ) ;
150+
151+ // Create the HTTP agent
152+ let agent = HttpAgent :: new ( base_url, headers) ;
153+
154+ let subscriber = GenerativeUiSubscriber :: new ( ) ;
155+
156+ // Create run parameters for testing generative UI with planning
157+ let params = RunAgentParams {
158+ messages : vec ! [
159+ Message :: User {
160+ id: MessageId :: random( ) ,
161+ content: "I need to organize a birthday party for my friend. Can you help me \
162+ . into( ) ,
163+ name: None ,
164+ }
165+ ] ,
166+ forwarded_props : Some ( JsonValue :: Null ) ,
167+ ..Default :: default ( )
168+ } ;
169+
170+ info ! ( "Starting generative UI agent run..." ) ;
171+ info ! ( "Testing planning functionality with state snapshots and deltas" ) ;
172+
173+ let result = agent. run_agent ( & params, vec ! [ Arc :: new( subscriber) ] ) . await ?;
174+
175+ info ! ( "Agent run completed successfully!" ) ;
176+ info ! ( "Final result: {}" , result. result) ;
177+ info ! ( "Generated {} new messages" , result. new_messages. len( ) ) ;
178+ info ! ( "Final state: {:#?}" , result. new_state) ;
179+
180+ // Print the messages for debugging
181+ for ( i, message) in result. new_messages . iter ( ) . enumerate ( ) {
182+ info ! ( "Message {}: {:?}" , i + 1 , message) ;
183+ }
184+
185+ Ok ( ( ) )
186+ }
0 commit comments