1use schemars::JsonSchema;
2use serde::{Deserialize, Serialize};
3use serde_json::Value;
4
5#[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#[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 pub fn with_awp_version(mut self, version: impl Into<String>) -> Self {
63 self.awp_version = Some(version.into());
64 self
65 }
66
67 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 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}