adk_ui/tools/
surface_persistence.rs

1use crate::compat::{AdkError, Result, Tool, ToolContext};
2use crate::persistence::{SavedSurface, SurfaceStore};
3use crate::schema::{Component, UiUpdate};
4use crate::surface_runtime::{
5    next_surface_ref, observe_surface_version, record_surface_ref, surface_owner,
6};
7use async_trait::async_trait;
8use schemars::JsonSchema;
9use serde::{Deserialize, Serialize};
10use serde_json::{Value, json};
11use std::sync::Arc;
12
13fn store_error(error: impl std::fmt::Display) -> AdkError {
14    AdkError::tool(error.to_string())
15}
16
17#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
18pub struct SaveSurfaceParams {
19    pub id: String,
20    pub name: String,
21    pub payload: Value,
22    #[serde(default)]
23    pub expected_version: Option<u64>,
24}
25
26#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
27pub struct SurfaceIdParams {
28    pub id: String,
29}
30
31#[derive(Debug, Clone, Default, Serialize, Deserialize, JsonSchema)]
32pub struct ListSurfacesParams {}
33
34#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
35pub struct DeleteSurfaceParams {
36    pub id: String,
37    #[serde(default)]
38    pub expected_version: Option<u64>,
39}
40
41macro_rules! store_tool {
42    ($name:ident) => {
43        pub struct $name {
44            store: Arc<dyn SurfaceStore>,
45        }
46        impl $name {
47            pub fn new(store: Arc<dyn SurfaceStore>) -> Self {
48                Self { store }
49            }
50        }
51    };
52}
53
54store_tool!(SaveSurfaceTool);
55store_tool!(LoadSurfaceTool);
56store_tool!(ListSurfacesTool);
57store_tool!(DeleteSurfaceTool);
58
59#[async_trait]
60impl Tool for SaveSurfaceTool {
61    fn name(&self) -> &str {
62        "save_surface"
63    }
64    fn description(&self) -> &str {
65        "Save a named UI surface. Use expected_version to prevent overwriting a newer revision."
66    }
67    fn parameters_schema(&self) -> Option<Value> {
68        Some(super::generate_gemini_schema::<SaveSurfaceParams>())
69    }
70    async fn execute(&self, ctx: Arc<dyn ToolContext>, args: Value) -> Result<Value> {
71        let params: SaveSurfaceParams = serde_json::from_value(args).map_err(store_error)?;
72        let owner = surface_owner(ctx.as_ref());
73        let saved = self
74            .store
75            .save(
76                &owner,
77                &params.id,
78                &params.name,
79                params.payload,
80                params.expected_version,
81            )
82            .await
83            .map_err(store_error)?;
84        let surface_ref = observe_surface_version(&ctx, saved.id.clone(), saved.version);
85        record_surface_ref(&ctx, &surface_ref);
86        serde_json::to_value(saved).map_err(store_error)
87    }
88}
89
90#[async_trait]
91impl Tool for LoadSurfaceTool {
92    fn name(&self) -> &str {
93        "load_surface"
94    }
95    fn description(&self) -> &str {
96        "Load a saved UI surface owned by the current agent."
97    }
98    fn parameters_schema(&self) -> Option<Value> {
99        Some(super::generate_gemini_schema::<SurfaceIdParams>())
100    }
101    async fn execute(&self, ctx: Arc<dyn ToolContext>, args: Value) -> Result<Value> {
102        let params: SurfaceIdParams = serde_json::from_value(args).map_err(store_error)?;
103        let saved = self
104            .store
105            .load(&surface_owner(ctx.as_ref()), &params.id)
106            .await
107            .map_err(store_error)?;
108        let surface_ref = observe_surface_version(&ctx, saved.id.clone(), saved.version);
109        record_surface_ref(&ctx, &surface_ref);
110        serde_json::to_value(saved).map_err(store_error)
111    }
112}
113
114#[async_trait]
115impl Tool for ListSurfacesTool {
116    fn name(&self) -> &str {
117        "list_surfaces"
118    }
119    fn description(&self) -> &str {
120        "List saved UI surfaces owned by the current agent."
121    }
122    fn parameters_schema(&self) -> Option<Value> {
123        Some(super::generate_gemini_schema::<ListSurfacesParams>())
124    }
125    async fn execute(&self, ctx: Arc<dyn ToolContext>, args: Value) -> Result<Value> {
126        let _: ListSurfacesParams = serde_json::from_value(args).map_err(store_error)?;
127        let surfaces = self
128            .store
129            .list(&surface_owner(ctx.as_ref()))
130            .await
131            .map_err(store_error)?;
132        serde_json::to_value(surfaces).map_err(store_error)
133    }
134}
135
136#[async_trait]
137impl Tool for DeleteSurfaceTool {
138    fn name(&self) -> &str {
139        "delete_surface"
140    }
141    fn description(&self) -> &str {
142        "Delete a saved UI surface. Use expected_version to prevent deleting a newer revision."
143    }
144    fn parameters_schema(&self) -> Option<Value> {
145        Some(super::generate_gemini_schema::<DeleteSurfaceParams>())
146    }
147    async fn execute(&self, ctx: Arc<dyn ToolContext>, args: Value) -> Result<Value> {
148        let params: DeleteSurfaceParams = serde_json::from_value(args).map_err(store_error)?;
149        let deleted = self
150            .store
151            .delete(
152                &surface_owner(ctx.as_ref()),
153                &params.id,
154                params.expected_version,
155            )
156            .await
157            .map_err(store_error)?;
158        let mut actions = ctx.actions();
159        actions.state_delta.insert(
160            "adk_ui.deleted_surface".to_string(),
161            json!({ "id": params.id, "deleted": deleted }),
162        );
163        ctx.set_actions(actions);
164        Ok(json!({ "id": params.id, "deleted": deleted }))
165    }
166}
167
168#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
169pub struct PatchSurfaceParams {
170    #[serde(default = "default_surface_id")]
171    pub surface_id: String,
172    pub target_id: String,
173    pub component: Component,
174    #[serde(default)]
175    pub expected_version: Option<u64>,
176}
177
178fn default_surface_id() -> String {
179    "main".to_string()
180}
181
182pub struct PatchSurfaceTool {
183    store: Option<Arc<dyn SurfaceStore>>,
184}
185
186impl PatchSurfaceTool {
187    pub fn new() -> Self {
188        Self { store: None }
189    }
190    pub fn with_store(store: Arc<dyn SurfaceStore>) -> Self {
191        Self { store: Some(store) }
192    }
193}
194
195impl Default for PatchSurfaceTool {
196    fn default() -> Self {
197        Self::new()
198    }
199}
200
201fn merge_value(target: &mut Value, patch: &Value) {
202    match (target, patch) {
203        (Value::Object(target), Value::Object(patch)) => {
204            for (key, value) in patch {
205                target.insert(key.clone(), value.clone());
206            }
207        }
208        (target, patch) => *target = patch.clone(),
209    }
210}
211
212fn patch_component(value: &mut Value, target_id: &str, patch: &Value) -> bool {
213    match value {
214        Value::Object(object) => {
215            if object.get("id").and_then(Value::as_str) == Some(target_id) {
216                merge_value(value, patch);
217                return true;
218            }
219            object
220                .values_mut()
221                .any(|child| patch_component(child, target_id, patch))
222        }
223        Value::Array(values) => values
224            .iter_mut()
225            .any(|child| patch_component(child, target_id, patch)),
226        _ => false,
227    }
228}
229
230#[async_trait]
231impl Tool for PatchSurfaceTool {
232    fn name(&self) -> &str {
233        "patch_surface"
234    }
235    fn description(&self) -> &str {
236        "Patch a component on an existing surface and emit an incremental UiUpdate. When persistence is configured, the saved surface is updated atomically."
237    }
238    fn parameters_schema(&self) -> Option<Value> {
239        Some(super::generate_gemini_schema::<PatchSurfaceParams>())
240    }
241    async fn execute(&self, ctx: Arc<dyn ToolContext>, args: Value) -> Result<Value> {
242        let params: PatchSurfaceParams = serde_json::from_value(args).map_err(store_error)?;
243        let update = UiUpdate::patch(params.target_id.clone(), params.component.clone());
244        let surface_ref = if let Some(store) = &self.store {
245            let owner = surface_owner(ctx.as_ref());
246            let current = store
247                .load(&owner, &params.surface_id)
248                .await
249                .map_err(store_error)?;
250            let mut payload = current.payload.clone();
251            let patch = serde_json::to_value(&params.component).map_err(store_error)?;
252            if !patch_component(&mut payload, &params.target_id, &patch) {
253                return Err(AdkError::tool(format!(
254                    "component not found: {}",
255                    params.target_id
256                )));
257            }
258            let expected = params.expected_version.or(Some(current.version));
259            let saved: SavedSurface = store
260                .save(&owner, &current.id, &current.name, payload, expected)
261                .await
262                .map_err(store_error)?;
263            observe_surface_version(&ctx, saved.id, saved.version)
264        } else {
265            next_surface_ref(&ctx, params.surface_id)
266        };
267        record_surface_ref(&ctx, &surface_ref);
268        Ok(json!({ "surface_ref": surface_ref, "update": update }))
269    }
270}
adk_ui/tools/surface_persistence.rs

adk_ui/tools/
surface_persistence.rs
