diff --git a/cloud/agent/agent.go b/cloud/agent/agent.go
index c5873eb..028a7d6 100644
--- a/cloud/agent/agent.go
+++ b/cloud/agent/agent.go
@@ -128,7 +128,7 @@
 	b.MaxElapsedTime = 0
 	// Main heartbeat loop
 	for {
-		req := bpb.AgentHeartbeatRequest{
+		req := bpb.HeartbeatRequest{
 			MachineId: agentInit.TakeoverInit.MachineId,
 		}
 		if sentFirstHeartBeat && !hwReportSent {
diff --git a/cloud/agent/e2e/main_test.go b/cloud/agent/e2e/main_test.go
index 84fd611..bcc38f3 100644
--- a/cloud/agent/e2e/main_test.go
+++ b/cloud/agent/e2e/main_test.go
@@ -67,8 +67,8 @@
 	installationReport  *bpb.OSInstallationReport
 }
 
-func (f *fakeServer) Heartbeat(ctx context.Context, req *bpb.AgentHeartbeatRequest) (*bpb.AgentHeartbeatResponse, error) {
-	var res bpb.AgentHeartbeatResponse
+func (f *fakeServer) Heartbeat(ctx context.Context, req *bpb.HeartbeatRequest) (*bpb.HeartbeatResponse, error) {
+	var res bpb.HeartbeatResponse
 	if req.HardwareReport != nil {
 		f.hardwareReport = req.HardwareReport
 	}