diff --git a/code/game/objects/items/robot/ai_upgrades.dm b/code/game/objects/items/robot/ai_upgrades.dm
index 39aff96c62c..b434ad34949 100644
--- a/code/game/objects/items/robot/ai_upgrades.dm
+++ b/code/game/objects/items/robot/ai_upgrades.dm
@@ -9,12 +9,12 @@
 	icon_state = "datadisk3"
 
 
-/obj/item/malf_upgrade/afterattack(mob/living/silicon/ai/AI, mob/user, proximity)
-	. = ..()
+/obj/item/malf_upgrade/pre_attack(atom/A, mob/living/user, proximity)
 	if(!proximity)
-		return
-	if(!istype(AI))
-		return
+		return ..()
+	if(!isAI(A))
+		return ..()
+	var/mob/living/silicon/ai/AI = A
 	if(AI.malf_picker)
 		AI.malf_picker.processing_time += 50
 		to_chat(AI, "<span class='userdanger'>[user] has attempted to upgrade you with combat software that you already possess. You gain 50 points to spend on Malfunction Modules instead.</span>")
@@ -27,6 +27,7 @@
 		message_admins("[ADMIN_LOOKUPFLW(user)] has upgraded [ADMIN_LOOKUPFLW(AI)] with a [src].")
 	to_chat(user, "<span class='notice'>You upgrade [AI]. [src] is consumed in the process.</span>")
 	qdel(src)
+	return TRUE
 
 
 //Lipreading
@@ -36,12 +37,12 @@
 	icon = 'icons/obj/module.dmi'
 	icon_state = "datadisk3"
 
-/obj/item/surveillance_upgrade/afterattack(mob/living/silicon/ai/AI, mob/user, proximity)
-	. = ..()
+/obj/item/surveillance_upgrade/pre_attack(atom/A, mob/living/user, proximity)
 	if(!proximity)
-		return
-	if(!istype(AI))
-		return
+		return ..()
+	if(!isAI(A))
+		return ..()
+	var/mob/living/silicon/ai/AI = A
 	if(AI.eyeobj)
 		AI.eyeobj.relay_speech = TRUE
 		to_chat(AI, "<span class='userdanger'>[user] has upgraded you with surveillance software!</span>")
@@ -50,3 +51,4 @@
 	log_game("[key_name(user)] has upgraded [key_name(AI)] with a [src].")
 	message_admins("[ADMIN_LOOKUPFLW(user)] has upgraded [ADMIN_LOOKUPFLW(AI)] with a [src].")
 	qdel(src)
+	return TRUE