1use crate::*;
2use rmcp::{
3 ErrorData, RoleServer, ServerHandler,
4 model::{
5 CacheScope, CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock,
6 ListToolsResult, PaginatedRequestParams, ServerCapabilities, ServerConfig, Tool,
7 ToolAnnotations,
8 },
9 service::RequestContext,
10};
11
12#[derive(Clone)]
13struct Observability(Arc<App>);
14impl ServerHandler for Observability {
15 fn get_info(&self) -> ServerConfig {
16 ServerConfig::new(ServerCapabilities::builder().enable_tools().build())
17 }
18 async fn list_tools(
19 &self,
20 _: Option<PaginatedRequestParams>,
21 _: RequestContext<RoleServer>,
22 ) -> std::result::Result<ListToolsResult, ErrorData> {
23 Ok(ListToolsResult { tools: [
24 ("get_logs", "Retrieve logs from one granted service."),
25 ("get_traces", "Find traces from one granted service."),
26 ("get_trace", "Retrieve a trace with spans limited to granted services."),
27 ].into_iter().map(|(name, description)| {
28 let mut properties = json!({"service":{"type":"string"},"q":{"type":"string"},"limit":{"type":"integer","minimum":1,"maximum":500}});
29 let required = if name=="get_trace" { properties=json!({"service":{"type":"string"},"trace_id":{"type":"string"}}); json!(["service","trace_id"]) } else {json!(["service"])};
30 Tool::new(name, description, json!({"type":"object","properties":properties,"required":required,"additionalProperties":false}).as_object().unwrap().clone()).with_annotations(ToolAnnotations::new().read_only(true))
31 }).collect(), ttl_ms: Some(0), cache_scope: Some(CacheScope::Private), ..Default::default() })
32 }
33 async fn call_tool(
34 &self,
35 request: CallToolRequestParams,
36 context: RequestContext<RoleServer>,
37 ) -> std::result::Result<CallToolResponse, ErrorData> {
38 let result: Result<Value> = async {
39 let grant = &context
40 .extensions
41 .get::<axum::http::request::Parts>()
42 .and_then(|p| p.extensions.get::<mcp::Grant>())
43 .ok_or_else(|| Error::new(401, "This connection expired. Connect again."))?
44 .0;
45 let arguments = request.arguments.unwrap_or_default();
46 let service = arguments
47 .get("service")
48 .and_then(Value::as_str)
49 .ok_or_else(|| Error::new(400, "Choose a granted service."))?;
50 if !array(&grant["resources"]).iter().any(|id| id == service) {
51 return Err(Error::new(
52 403,
53 "This service is outside the connection's access.",
54 ));
55 }
56 let mut query = HashMap::new();
57 for (key, value) in &arguments {
58 match key.as_str() {
59 "service" => {}
60 "q" if request.name != "get_trace" => {
61 query.insert(
62 key.clone(),
63 value
64 .as_str()
65 .filter(|s| s.len() <= 4096)
66 .ok_or_else(|| Error::new(400, "Narrow the search."))?
67 .to_owned(),
68 );
69 }
70 "limit" if request.name != "get_trace" => {
71 query.insert(
72 key.clone(),
73 value
74 .as_u64()
75 .filter(|n| (1..=500).contains(n))
76 .ok_or_else(|| Error::new(400, "Choose a limit from 1 to 500."))?
77 .to_string(),
78 );
79 }
80 "trace_id" if request.name == "get_trace" && value.as_str().is_some() => {}
81 _ => return Err(Error::new(400, "Use the fields listed for this tool.")),
82 }
83 }
84 match request.name.as_ref() {
85 "get_logs" => {
86 Ok(json!({"logs":telemetry::logs(self.0.clone(), service, &query).await?}))
87 }
88 "get_traces" => {
89 let jobs = core::scan(self.0.clone()).await?;
90 let traces = telemetry::traces(self.0.clone(), service, &query, |span| {
91 visible_span(span, &jobs.value, &grant["resources"])
92 })
93 .await?;
94 Ok(json!({"traces":traces}))
95 }
96 "get_trace" => {
97 let id = arguments
98 .get("trace_id")
99 .and_then(Value::as_str)
100 .filter(|s| !s.is_empty() && s.len() <= 128)
101 .ok_or_else(|| Error::new(400, "Choose a trace ID."))?;
102 let mut trace = telemetry::trace(self.0.clone(), id).await?;
103 let jobs = core::scan(self.0.clone()).await?;
104 if !array(&trace["spans"])
105 .iter()
106 .any(|span| visible_span(span, &jobs.value, &json!([service])))
107 {
108 return Err(Error::new(404, "No trace from that service has this ID."));
109 }
110 trace["spans"]
111 .as_array_mut()
112 .unwrap()
113 .retain(|span| visible_span(span, &jobs.value, &grant["resources"]));
114 Ok(json!({"trace":trace}))
115 }
116 _ => Err(Error::new(404, "No tool with that name.")),
117 }
118 }
119 .await;
120 Ok(match result {
121 Ok(value) => CallToolResult::structured(value),
122 Err(error) => CallToolResult::error(vec![ContentBlock::text(if error.status >= 500 {
123 "The service couldn't answer. Check its dashboard and retry.".to_owned()
124 } else {
125 error.message
126 })]),
127 }
128 .into())
129 }
130}
131fn visible_span(span: &Value, jobs: &Value, resources: &Value) -> bool {
132 if let Some(id) = span["attributes"]["studio.service"].as_str() {
133 return array(resources).iter().any(|r| r == id);
134 }
135 let owners: Vec<_> = jobs
136 .as_object()
137 .into_iter()
138 .flat_map(|jobs| jobs.iter())
139 .filter(|(_, job)| job["job"]["Meta"]["studio_trace_service"] == span["service"])
140 .map(|(id, _)| id)
141 .collect();
142 owners.len() == 1 && array(resources).iter().any(|r| r == owners[0])
143}
144pub fn router(app: Arc<App>) -> Router {
145 let state = app.clone();
146 mcp::router(app, "observability", move || {
147 Ok(Observability(state.clone()))
148 })
149}
150
151#[cfg(test)]
152mod tests {
153 use super::*;
154 #[test]
155 fn trace_grants_reject_foreign_and_ambiguous_spans() {
156 let jobs = json!({"allowed":{"job":{"Meta":{"studio_trace_service":"backend"}}},"private":{"job":{"Meta":{"studio_trace_service":"private-api"}}}});
157 let resources = json!(["allowed"]);
158 assert!(visible_span(
159 &json!({"service":"backend","attributes":{}}),
160 &jobs,
161 &resources
162 ));
163 assert!(!visible_span(
164 &json!({"service":"private-api","attributes":{}}),
165 &jobs,
166 &resources
167 ));
168 assert!(!visible_span(
169 &json!({"service":"backend","attributes":{"studio.service":"private"}}),
170 &jobs,
171 &resources
172 ));
173 assert!(visible_span(
174 &json!({"service":"allowed","attributes":{"studio.service":"allowed"}}),
175 &jobs,
176 &resources
177 ));
178 let mut ambiguous = jobs.clone();
179 ambiguous["private"]["job"]["Meta"]["studio_trace_service"] = json!("backend");
180 assert!(!visible_span(
181 &json!({"service":"backend","attributes":{}}),
182 &ambiguous,
183 &resources
184 ));
185 assert!(!visible_span(
186 &json!({"service":"unknown","attributes":{}}),
187 &jobs,
188 &resources
189 ));
190 }
191}