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