Skip to main content

adk_ui/
surface_runtime.rs

1use crate::compat::ToolContext;
2use schemars::JsonSchema;
3use serde::{Deserialize, Serialize};
4use serde_json::{Value, json};
5use std::collections::HashMap;
6use std::sync::{Arc, Mutex, OnceLock};
7
8pub const ACTIVE_SURFACE_STATE_KEY: &str = "adk_ui.active_surface";
9pub const SURFACE_REFS_STATE_KEY: &str = "adk_ui.surface_refs";
10
11#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
12pub struct SurfaceRef {
13    pub id: String,
14    pub version: u64,
15}
16
17impl SurfaceRef {
18    pub fn new(id: impl Into<String>, version: u64) -> Self {
19        Self {
20            id: id.into(),
21            version,
22        }
23    }
24}
25
26fn revisions() -> &'static Mutex<HashMap<String, u64>> {
27    static REVISIONS: OnceLock<Mutex<HashMap<String, u64>>> = OnceLock::new();
28    REVISIONS.get_or_init(|| Mutex::new(HashMap::new()))
29}
30
31fn revision_key(ctx: &dyn ToolContext, surface_id: &str) -> String {
32    format!(
33        "{}\u{1f}{}\u{1f}{}\u{1f}{}",
34        ctx.app_name(),
35        ctx.user_id(),
36        ctx.session_id(),
37        surface_id
38    )
39}
40
41pub(crate) fn next_surface_ref(
42    ctx: &Arc<dyn ToolContext>,
43    surface_id: impl Into<String>,
44) -> SurfaceRef {
45    let surface_id = surface_id.into();
46    let key = revision_key(ctx.as_ref(), &surface_id);
47    let mut values = revisions()
48        .lock()
49        .unwrap_or_else(|poisoned| poisoned.into_inner());
50    let version = values
51        .entry(key)
52        .and_modify(|value| *value += 1)
53        .or_insert(1);
54    SurfaceRef::new(surface_id, *version)
55}
56
57pub(crate) fn observe_surface_version(
58    ctx: &Arc<dyn ToolContext>,
59    surface_id: impl Into<String>,
60    version: u64,
61) -> SurfaceRef {
62    let surface_id = surface_id.into();
63    let key = revision_key(ctx.as_ref(), &surface_id);
64    let mut values = revisions()
65        .lock()
66        .unwrap_or_else(|poisoned| poisoned.into_inner());
67    values
68        .entry(key)
69        .and_modify(|current| *current = (*current).max(version))
70        .or_insert(version);
71    SurfaceRef::new(surface_id, version)
72}
73
74pub(crate) fn record_surface_ref(ctx: &Arc<dyn ToolContext>, surface_ref: &SurfaceRef) {
75    let mut actions = ctx.actions();
76    actions.state_delta.insert(
77        ACTIVE_SURFACE_STATE_KEY.to_string(),
78        serde_json::to_value(surface_ref).unwrap_or_else(|_| {
79            json!({
80                "id": surface_ref.id,
81                "version": surface_ref.version,
82            })
83        }),
84    );
85    ctx.set_actions(actions);
86}
87
88pub(crate) fn record_surface_refs(
89    ctx: &Arc<dyn ToolContext>,
90    active: &SurfaceRef,
91    refs: &[SurfaceRef],
92) {
93    let mut actions = ctx.actions();
94    actions.state_delta.insert(
95        ACTIVE_SURFACE_STATE_KEY.to_string(),
96        serde_json::to_value(active).unwrap_or(Value::Null),
97    );
98    actions.state_delta.insert(
99        SURFACE_REFS_STATE_KEY.to_string(),
100        serde_json::to_value(refs).unwrap_or(Value::Array(vec![])),
101    );
102    ctx.set_actions(actions);
103}
104
105pub(crate) fn surface_owner(ctx: &dyn ToolContext) -> String {
106    let agent = ctx.agent_name().trim();
107    if !agent.is_empty() {
108        agent.to_string()
109    } else {
110        ctx.user_id().to_string()
111    }
112}