diff --git a/cloud/bmaas/server/agent_callback_service.go b/cloud/bmaas/server/agent_callback_service.go
index cbd912a..b664231 100644
--- a/cloud/bmaas/server/agent_callback_service.go
+++ b/cloud/bmaas/server/agent_callback_service.go
@@ -31,7 +31,7 @@
 	errAgentUnauthenticated = errors.New("machine id or public key unknown")
 )
 
-func (a *agentCallbackService) Heartbeat(ctx context.Context, req *apb.AgentHeartbeatRequest) (*apb.AgentHeartbeatResponse, error) {
+func (a *agentCallbackService) Heartbeat(ctx context.Context, req *apb.HeartbeatRequest) (*apb.HeartbeatResponse, error) {
 	// Extract ED25519 self-signed certificate from client connection.
 	cert, err := rpc.GetPeerCertificate(ctx)
 	if err != nil {
@@ -160,7 +160,7 @@
 		klog.Errorf("Failure during OS installation request retrieval: %v", err)
 	}
 
-	return &apb.AgentHeartbeatResponse{
+	return &apb.HeartbeatResponse{
 		InstallationRequest: installRequest,
 	}, nil
 }