extend instrumentation

Rhys Sullivan committed 6 months ago

Commit 3f791571 parent 88ad3581 file · +49 −10Browse files

Unified · Split

 export const McpAuthLive = Layer.succeed(McpAuth, {
   verifyBearer: (request) =>
-    Effect.promise(async () => {
+    Effect.gen(function* () {
       const authHeader = request.headers.get("authorization");
-      if (!authHeader?.startsWith(BEARER_PREFIX)) return null;
+      if (!authHeader?.startsWith(BEARER_PREFIX)) {
+        yield* Effect.annotateCurrentSpan({ "mcp.auth.outcome": "missing_bearer" });
+        return null;
+      }
       try {
-        const { payload } = await jwtVerify(authHeader.slice(BEARER_PREFIX.length), jwks, {
+        const { payload } = yield* Effect.promise(() =>
-          issuer: AUTHKIT_DOMAIN,
+          jwtVerify(authHeader.slice(BEARER_PREFIX.length), jwks, {
-        });
+            issuer: AUTHKIT_DOMAIN,
-        if (!payload.sub) return null;
+          }),
-        return {
+        ).pipe(Effect.withSpan("mcp.auth.jwt_verify"));
+        if (!payload.sub) {
+          yield* Effect.annotateCurrentSpan({ "mcp.auth.outcome": "missing_subject" });
+          return null;
+        }
+        const token = {
           accountId: payload.sub,
           organizationId: (payload.org_id as string | undefined) ?? null,
         };
+        yield* Effect.annotateCurrentSpan({
+          "mcp.auth.outcome": "verified",
+          "mcp.auth.has_organization": !!token.organizationId,
+        });
+        return token;
       } catch {
+        yield* Effect.annotateCurrentSpan({ "mcp.auth.outcome": "invalid" });
         return null;
       }
-    }),
+    }).pipe(Effect.withSpan("mcp.auth.verify_bearer")),
 });
 // ---------------------------------------------------------------------------
     } catch {
       return Option.none();
     }
-  });
+  }).pipe(Effect.withSpan("mcp.request.read_json_rpc"));
 const methodAttrs = (envelope: JsonRpcEnvelope): Record<string, unknown> => {
   const params = envelope.params ?? {};
       ...baseAttrs,
       ...rpcAttrs(envelope),
     });
-  });
+  }).pipe(
+    Effect.withSpan("mcp.request.annotate", {
+      attributes: {
+        "mcp.request.method": request.method,
+        "mcp.request.parse_body": opts.parseBody,
+      },
+    }),
+  );
 // ---------------------------------------------------------------------------
 // OAuth metadata endpoints
     const propagated = withPropagationHeaders(request, propagation);
     const raw = yield* Effect.promise(
       () => stub.handleRequest(propagated) as Promise<Response>,
+    ).pipe(
+      Effect.withSpan("mcp.do.handle_request", {
+        attributes: {
+          "mcp.request.method": request.method,
+          "mcp.request.session_id_present": true,
+        },
+      }),
     );
     const annotated = peek ? yield* peekAndAnnotate(raw) : raw;
     return HttpServerResponse.raw(annotated);
     const propagation = yield* currentPropagationHeaders(request);
     yield* Effect.promise(() =>
       stub.init({ organizationId, userId: token.accountId }, propagation),
+    ).pipe(
+      Effect.withSpan("mcp.do.init", {
+        attributes: { "mcp.request.session_id_present": false },
+      }),
     );
     const propagated = withPropagationHeaders(request, propagation);
     const raw = yield* Effect.promise(
       () => stub.handleRequest(propagated) as Promise<Response>,
+    ).pipe(
+      Effect.withSpan("mcp.do.handle_request", {
+        attributes: {
+          "mcp.request.method": request.method,
+          "mcp.request.session_id_present": false,
+        },
+      }),
     );
     const annotated = yield* peekAndAnnotate(raw);
     return HttpServerResponse.raw(annotated);