Trace Auth with status_code and reason

Changeset

Unit tests for auth.rs

Add raw_operation as a telemetry attribute

Unit test starting a streamable http server

Rename self to operation_source

Author: swcollard

.changesets/feat_telemetry_operations_trace_metrics.md

@@ -0,0 +1,4 @@
+### Telemetry: Trace operations and auth - @swcollard PR #375
+
+* Adds traces for the MCP server generating Tools from Operations and performing authorization
+* Includes the HTTP status code to the top level HTTP trace

crates/apollo-mcp-server/src/auth.rs

@@ -84,6 +84,7 @@ impl Config {
 }
 
 /// Validate that requests made have a corresponding bearer JWT token
+#[tracing::instrument(skip_all, fields(status_code, reason))]
 async fn oauth_validate(
     State(auth_config): State<Config>,
     token: Option<TypedHeader<Authorization<Bearer>>>,
@@ -104,17 +105,85 @@ async fn oauth_validate(
     };
 
     let validator = NetworkedTokenValidator::new(&auth_config.audiences, &auth_config.servers);
-    let token = token.ok_or_else(unauthorized_error)?;
-
-    let valid_token = validator
-        .validate(token.0)
-        .await
-        .ok_or_else(unauthorized_error)?;
+    let token = token.ok_or_else(|| {
+        tracing::Span::current().record("reason", "missing_token");
+        tracing::Span::current().record("status_code", StatusCode::UNAUTHORIZED.as_u16());
+        unauthorized_error()
+    })?;
+
+    let valid_token = validator.validate(token.0).await.ok_or_else(|| {
+        tracing::Span::current().record("reason", "invalid_token");
+        tracing::Span::current().record("status_code", StatusCode::UNAUTHORIZED.as_u16());
+        unauthorized_error()
+    })?;
 
     // Insert new context to ensure that handlers only use our enforced token verification
     // for propagation
     request.extensions_mut().insert(valid_token);
 
     let response = next.run(request).await;
+    tracing::Span::current().record("status_code", response.status().as_u16());
     Ok(response)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use axum::middleware::from_fn_with_state;
+    use axum::routing::get;
+    use axum::{
+        Router,
+        body::Body,
+        http::{Request, StatusCode},
+    };
+    use http::header::{AUTHORIZATION, WWW_AUTHENTICATE};
+    use tower::ServiceExt; // for .oneshot()
+    use url::Url;
+
+    fn test_config() -> Config {
+        Config {
+            servers: vec![Url::parse("http://localhost:1234").unwrap()],
+            audiences: vec!["test-audience".to_string()],
+            resource: Url::parse("http://localhost:4000").unwrap(),
+            resource_documentation: None,
+            scopes: vec!["read".to_string()],
+            disable_auth_token_passthrough: false,
+        }
+    }
+
+    fn test_router(config: Config) -> Router {
+        Router::new()
+            .route("/test", get(|| async { "ok" }))
+            .layer(from_fn_with_state(config, oauth_validate))
+    }
+
+    #[tokio::test]
+    async fn missing_token_returns_unauthorized() {
+        let config = test_config();
+        let app = test_router(config.clone());
+        let req = Request::builder().uri("/test").body(Body::empty()).unwrap();
+        let res = app.oneshot(req).await.unwrap();
+        assert_eq!(res.status(), StatusCode::UNAUTHORIZED);
+        let headers = res.headers();
+        let www_auth = headers.get(WWW_AUTHENTICATE).unwrap().to_str().unwrap();
+        assert!(www_auth.contains("Bearer"));
+        assert!(www_auth.contains("resource_metadata"));
+    }
+
+    #[tokio::test]
+    async fn invalid_token_returns_unauthorized() {
+        let config = test_config();
+        let app = test_router(config.clone());
+        let req = Request::builder()
+            .uri("/test")
+            .header(AUTHORIZATION, "Bearer invalidtoken")
+            .body(Body::empty())
+            .unwrap();
+        let res = app.oneshot(req).await.unwrap();
+        assert_eq!(res.status(), StatusCode::UNAUTHORIZED);
+        let headers = res.headers();
+        let www_auth = headers.get(WWW_AUTHENTICATE).unwrap().to_str().unwrap();
+        assert!(www_auth.contains("Bearer"));
+        assert!(www_auth.contains("resource_metadata"));
+    }
+}

crates/apollo-mcp-server/src/operations/operation_source.rs

@@ -38,7 +38,7 @@ pub enum OperationSource {
 }
 
 impl OperationSource {
-    #[tracing::instrument]
+    #[tracing::instrument(skip_all, fields(operation_source = ?self))]
     pub async fn into_stream(self) -> impl Stream<Item = Event> {
         match self {
             OperationSource::Files(paths) => Self::stream_file_changes(paths).boxed(),

crates/apollo-mcp-server/src/server/states/running.rs

@@ -105,6 +105,7 @@ impl Running {
         Ok(self)
     }
 
+    #[tracing::instrument(skip_all)]
     pub(super) async fn update_operations(
         self,
         operations: Vec<RawOperation>,
@@ -146,6 +147,7 @@ impl Running {
     }
 
     /// Notify any peers that tools have changed. Drops unreachable peers from the list.
+    #[tracing::instrument(skip_all)]
     async fn notify_tool_list_changed(peers: Arc<RwLock<Vec<Peer<RoleServer>>>>) {
         let mut peers = peers.write().await;
         if !peers.is_empty() {

crates/apollo-mcp-server/src/server/states/starting.rs

@@ -229,10 +229,7 @@ impl Starting {
                             |response: &axum::http::Response<_>,
                              _latency: std::time::Duration,
                              span: &tracing::Span| {
-                                span.record(
-                                    "status_code",
-                                    tracing::field::display(response.status()),
-                                );
+                                span.record("status", tracing::field::display(response.status()));
                             },
                         ),
                 );
@@ -334,3 +331,49 @@ async fn health_endpoint(
 
     Ok((status_code, Json(json!(health))))
 }
+
+#[cfg(test)]
+mod tests {
+    use http::HeaderMap;
+    use url::Url;
+
+    use crate::health::HealthCheckConfig;
+
+    use super::*;
+
+    #[tokio::test]
+    async fn start_basic_server() {
+        let starting = Starting {
+            config: Config {
+                transport: Transport::StreamableHttp {
+                    auth: None,
+                    address: "127.0.0.1".parse().unwrap(),
+                    port: 7799,
+                    stateful_mode: false,
+                },
+                endpoint: Url::parse("http://localhost:4000").expect("valid url"),
+                mutation_mode: MutationMode::All,
+                execute_introspection: true,
+                headers: HeaderMap::new(),
+                validate_introspection: true,
+                introspect_introspection: true,
+                search_introspection: true,
+                introspect_minify: false,
+                search_minify: false,
+                explorer_graph_ref: None,
+                custom_scalar_map: None,
+                disable_type_description: false,
+                disable_schema_description: false,
+                disable_auth_token_passthrough: false,
+                search_leaf_depth: 5,
+                index_memory_bytes: 1024 * 1024 * 1024,
+                health_check: HealthCheckConfig::default(),
+            },
+            schema: Schema::parse_and_validate("type Query { hello: String }", "test.graphql")
+                .expect("Valid schema"),
+            operations: vec![],
+        };
+        let running = starting.start();
+        assert!(running.await.is_ok());
+    }
+}

crates/apollo-mcp-server/src/telemetry_attributes.rs

@@ -20,6 +20,9 @@ impl TelemetryAttribute {
             TelemetryAttribute::RequestId => {
                 Key::from_static_str(TelemetryAttribute::RequestId.as_str())
             }
+            TelemetryAttribute::RawOperation => {
+                Key::from_static_str(TelemetryAttribute::RawOperation.as_str())
+            }
         }
     }
 

crates/apollo-mcp-server/telemetry.toml

@@ -4,6 +4,7 @@ operation_id = "The operation id - either persisted query id, operation name, or
 operation_source = "The operation source - either operation (local file/op collection), persisted query, or LLM generated"
 request_id = "The request id"
 success = "Sucess flag indicator"
+raw_operation = "Graphql operation text and metadata used for Tool generation"
 
 [metrics.apollo.mcp]
 "initialize.count" = "Number of times initialize has been called"