feat: Agent Ping

agent/client/client.go

@@ -61,6 +61,8 @@ func (c *Client) Start(ctx context.Context) error {
 		return err
 	}
 
+	c.startHearthBeat(ctx)
+
 	return nil
 }
 

agent/client/workflow_ping.go

@@ -0,0 +1,23 @@
+package client
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/kubeshop/tracetest/agent/proto"
+)
+
+func (c *Client) startHearthBeat(ctx context.Context) error {
+	client := proto.NewOrchestratorClient(c.conn)
+	ticker := time.NewTicker(2 * time.Second)
+
+	go func() {
+		for range ticker.C {
+			fmt.Println("@@PING")
+			client.Ping(ctx, c.sessionConfig.AgentIdentification)
+		}
+	}()
+
+	return nil
+}

agent/proto/orchestrator.proto

@@ -25,6 +25,9 @@ service Orchestrator {
 
     // Register an agent to listen for shutdown commands
     rpc RegisterShutdownListener(AgentIdentification) returns (stream ShutdownRequest) {}
+
+    // Ping is used to check if the agent is still connected
+    rpc Ping(AgentIdentification) returns (Empty) {}
 }
 
 // Empty message for endpoints that don't return anything