Use reflection

Author: njooma

Makefile

@@ -57,6 +57,7 @@ build-buf: $(node_modules) clean-buf
 	$(buf) generate buf.build/googleapis/googleapis
 	$(buf) generate buf.build/viamrobotics/api:$$(cat api_version.lock) --path common,component,robot,service,app,provisioning,tagger,stream
 	$(buf) generate buf.build/viamrobotics/goutils
+	$(buf) generate buf.build/grpc/grpc --path grpc/reflection/v1/reflection.proto
 
 # js targets
 

examples/vanilla/package-lock.json

@@ -170,12 +170,15 @@ export class RobotClient extends EventDispatcher implements Robot {
         this.onDisconnect();
       }
     );
-    this.sessionManager = new SessionManager((): Transport => {
-      if (!this.transport) {
-        throw new Error(RobotClient.notConnectedYetStr);
+    this.sessionManager = new SessionManager(
+      this.serviceHost,
+      (): Transport => {
+        if (!this.transport) {
+          throw new Error(RobotClient.notConnectedYetStr);
+        }
+        return this.transport;
       }
-      return this.transport;
-    });
+    );
 
     // For each connection event type, add a listener to capture that
     // event and re-emit it with the 'connectionstatechange' event

src/robot/client.ts

@@ -27,7 +27,7 @@ const mockGetHeartBeatWindow = new Duration({
 
 describe('SessionManager', () => {
   beforeEach(() => {
-    sm = new SessionManager(() => mockTransport);
+    sm = new SessionManager('', () => mockTransport);
   });
 
   it('no session initially', () => {

src/robot/session-manager.spec.ts

@@ -1,9 +1,23 @@
+/* eslint-disable max-depth */
+import {
+  BinaryReader,
+  FileDescriptorProto,
+  MethodOptions,
+} from '@bufbuild/protobuf';
 import {
   Code,
   ConnectError,
   createClient,
   type Transport,
 } from '@connectrpc/connect';
+import { createAsyncIterable } from '@connectrpc/connect/protocol';
+import { safety_heartbeat_monitored as safteyHeartbeatMonitored } from '../gen/common/v1/common_pb';
+import { ServerReflection } from '../gen/grpc/reflection/v1/reflection_connect';
+import {
+  FileDescriptorResponse,
+  ListServiceResponse,
+  ServerReflectionRequest,
+} from '../gen/grpc/reflection/v1/reflection_pb';
 import { RobotService } from '../gen/robot/v1/robot_connect';
 import {
   SendSessionHeartbeatRequest,
@@ -22,29 +36,14 @@ const timeoutBlob = new Blob(
 );
 
 export default class SessionManager {
-  public readonly transport: Transport;
+  public static heartbeatMonitoredMethods: Record<string, boolean> = {};
 
-  public static readonly heartbeatMonitoredMethods = new Set<string>([
-    '/viam.component.arm.v1.ArmService/MoveToPosition',
-    '/viam.component.arm.v1.ArmService/MoveToJointPositions',
-    '/viam.component.arm.v1.ArmService/MoveThroughJointPositions',
-    '/viam.component.base.v1.BaseService/MoveStraight',
-    '/viam.component.base.v1.BaseService/Spin',
-    '/viam.component.base.v1.BaseService/SetPower',
-    '/viam.component.base.v1.BaseService/SetVelocity',
-    '/viam.component.gantry.v1.GantryService/MoveToPosition',
-    '/viam.component.gripper.v1.GripperService/Open',
-    '/viam.component.gripper.v1.GripperService/Grab',
-    '/viam.component.motor.v1.MotorService/SetPower',
-    '/viam.component.motor.v1.MotorService/GoFor',
-    '/viam.component.motor.v1.MotorService/GoTo',
-    '/viam.component.motor.v1.MotorService/SetRPM',
-    '/viam.component.servo.v1.ServoService/Move',
-  ]);
+  public readonly transport: Transport;
 
   private currentSessionID = '';
   private sessionsSupported: boolean | undefined;
   private heartbeatIntervalMs: number | undefined;
+  private host = '';
 
   private starting: Promise<void> | undefined;
 
@@ -53,7 +52,11 @@ export default class SessionManager {
     return createClient(RobotService, transport);
   }
 
-  constructor(private deferredTransport: () => Transport) {
+  constructor(
+    host: string,
+    private deferredTransport: () => Transport
+  ) {
+    this.host = host;
     this.transport = new SessionTransport(this.deferredTransport, this);
   }
 
@@ -185,6 +188,7 @@ export default class SessionManager {
           (Number(heartbeatWindow.seconds) * 1e3 +
             heartbeatWindow.nanos / 1e6) /
           5;
+        await this.applyHeartbeatMonitoredMethods();
         resolve();
         this.heartbeat().catch(console.error); // eslint-disable-line no-console
       })()
@@ -198,4 +202,59 @@ export default class SessionManager {
 
     return this.getSessionMetadataInner();
   }
+
+  private async applyHeartbeatMonitoredMethods(): Promise<void> {
+    const client = createClient(ServerReflection, this.transport);
+    const request = new ServerReflectionRequest({
+      host: this.host,
+      messageRequest: { case: 'listServices', value: '' },
+    });
+    const responseStream = client.serverReflectionInfo(
+      createAsyncIterable([request]),
+      { timeoutMs: 10_000 }
+    );
+    for await (const serviceResponse of responseStream) {
+      const fdpRequests = (
+        serviceResponse.messageResponse.value as ListServiceResponse
+      ).service.map((service) => {
+        return new ServerReflectionRequest({
+          messageRequest: { case: 'fileContainingSymbol', value: service.name },
+        });
+      });
+      const fdpResponseStream = client.serverReflectionInfo(
+        createAsyncIterable(fdpRequests),
+        { timeoutMs: 10_000 }
+      );
+      for await (const fdpResponse of fdpResponseStream) {
+        for (const fdp of (
+          fdpResponse.messageResponse.value as FileDescriptorResponse
+        ).fileDescriptorProto) {
+          const protoFile = FileDescriptorProto.fromBinary(fdp);
+          for (const service of protoFile.service) {
+            for (const method of service.method) {
+              SessionManager.heartbeatMonitoredMethods[
+                `/${protoFile.package}.${service.name}/${method.name}`
+              ] = SessionManager.hasHeartbeatOption(method.options);
+            }
+          }
+        }
+      }
+    }
+  }
+
+  private static hasHeartbeatOption(options?: MethodOptions): boolean {
+    if (!options) {
+      return false;
+    }
+    const reader = new BinaryReader(options.toBinary());
+    while (reader.pos < reader.len) {
+      const tag = reader.tag();
+      const [fieldNumber] = tag;
+      if (fieldNumber === safteyHeartbeatMonitored.field.no) {
+        return true;
+      }
+      reader.string();
+    }
+    return false;
+  }
 }

src/robot/session-manager.ts

@@ -20,7 +20,7 @@ export default class SessionTransport implements Transport {
   constructor(
     protected readonly deferredTransport: () => Transport,
     protected readonly sessionManager: SessionManager
-  ) {}
+  ) { }
 
   private async getSessionMetadata(): Promise<Headers> {
     try {
@@ -51,7 +51,7 @@ export default class SessionTransport implements Transport {
   ): Promise<UnaryResponse<I, O>> {
     const newHeaders = cloneHeaders(header);
     const methodPath = `/${service.typeName}/${method.name}`;
-    if (SessionManager.heartbeatMonitoredMethods.has(methodPath)) {
+    if (SessionManager.heartbeatMonitoredMethods[methodPath] ?? false) {
       const md = await this.getSessionMetadata();
       for (const [key, value] of md) {
         newHeaders.set(key, value);
@@ -82,7 +82,7 @@ export default class SessionTransport implements Transport {
   ): Promise<StreamResponse<I, O>> {
     const newHeaders = cloneHeaders(header);
     const methodPath = `/${service.typeName}/${method.name}`;
-    if (SessionManager.heartbeatMonitoredMethods.has(methodPath)) {
+    if (SessionManager.heartbeatMonitoredMethods[methodPath] ?? false) {
       const md = await this.getSessionMetadata();
       for (const [key, value] of md) {
         newHeaders.set(key, value);