Skip to main content

adk_ui/model/
envelope.rs

1use schemars::JsonSchema;
2use serde::{Deserialize, Serialize};
3use serde_json::Value;
4
5/// Standard envelope protocol marker for render tool outputs.
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
7#[serde(rename_all = "snake_case")]
8pub enum ToolEnvelopeProtocol {
9    AdkUi,
10    A2ui,
11    AgUi,
12    McpApps,
13    #[cfg(feature = "awp")]
14    Awp,
15}
16
17/// Canonical tool output envelope.
18///
19/// Payload fields are flattened to preserve backward-compatible response shapes
20/// while still attaching protocol metadata.
21#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
22pub struct ToolEnvelope<P> {
23    pub protocol: ToolEnvelopeProtocol,
24    pub version: String,
25    pub surface_id: String,
26    #[serde(flatten)]
27    pub payload: P,
28    #[serde(skip_serializing_if = "Option::is_none")]
29    pub meta: Option<Value>,
30    #[cfg(feature = "awp")]
31    #[serde(skip_serializing_if = "Option::is_none")]
32    pub awp_version: Option<String>,
33    #[cfg(feature = "awp")]
34    #[serde(skip_serializing_if = "Option::is_none")]
35    pub request_id: Option<String>,
36}
37
38impl<P> ToolEnvelope<P> {
39    pub fn new(protocol: ToolEnvelopeProtocol, surface_id: impl Into<String>, payload: P) -> Self {
40        Self {
41            protocol,
42            version: "1.0".to_string(),
43            surface_id: surface_id.into(),
44            payload,
45            meta: None,
46            #[cfg(feature = "awp")]
47            awp_version: None,
48            #[cfg(feature = "awp")]
49            request_id: None,
50        }
51    }
52
53    pub fn with_meta(mut self, meta: Option<Value>) -> Self {
54        self.meta = meta;
55        self
56    }
57}
58
59#[cfg(feature = "awp")]
60impl<P: Serialize> ToolEnvelope<P> {
61    /// Set the AWP version on this envelope.
62    pub fn with_awp_version(mut self, version: impl Into<String>) -> Self {
63        self.awp_version = Some(version.into());
64        self
65    }
66
67    /// Set the AWP request ID on this envelope.
68    pub fn with_request_id(mut self, request_id: impl Into<String>) -> Self {
69        self.request_id = Some(request_id.into());
70        self
71    }
72
73    /// Convert this envelope into an AWP response.
74    pub fn to_awp_response(&self) -> Result<awp_types::AwpResponse, crate::compat::AdkError> {
75        let payload = serde_json::to_value(&self.payload).map_err(|e| {
76            crate::compat::AdkError::tool(format!(
77                "Failed to serialize envelope payload for AWP response: {}",
78                e
79            ))
80        })?;
81        Ok(awp_types::AwpResponse {
82            id: uuid::Uuid::now_v7(),
83            version: awp_types::CURRENT_VERSION,
84            status: "ok".to_string(),
85            payload,
86        })
87    }
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93    use serde::Serialize;
94    use serde_json::json;
95
96    #[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
97    struct Payload {
98        value: String,
99    }
100
101    #[test]
102    fn envelope_serializes_flattened_payload() {
103        let envelope = ToolEnvelope::new(
104            ToolEnvelopeProtocol::A2ui,
105            "main",
106            Payload {
107                value: "ok".to_string(),
108            },
109        )
110        .with_meta(Some(json!({"trace_id": "abc"})));
111
112        let value = serde_json::to_value(envelope).expect("serialize");
113        assert_eq!(value["protocol"], "a2ui");
114        assert_eq!(value["surface_id"], "main");
115        assert_eq!(value["version"], "1.0");
116        assert_eq!(value["value"], "ok");
117        assert_eq!(value["meta"]["trace_id"], "abc");
118    }
119}