diff --git a/models/account_invoice.py b/models/account_invoice.py
index c34b24c..3f9c8d8 100644
--- a/models/account_invoice.py
+++ b/models/account_invoice.py
@@ -37,5 +37,7 @@ class AccountInvoice(models.Model):
             "context": {
                 "self_id": order.contract_annex_id.id,
                 "active_model": "res.partner.contract.annex",
+                "attachment_model": self._name,
+                "attachment_res_id": self.id
             },
         }
diff --git a/models/res_partner_contract_annex.py b/models/res_partner_contract_annex.py
index 6988cd1..15f7dd0 100644
--- a/models/res_partner_contract_annex.py
+++ b/models/res_partner_contract_annex.py
@@ -142,7 +142,7 @@ class ContractOrderAnnex(models.Model, IDocument, Extension):
                 "act_at": "{counter}.2 {type} {name}-2",
                 "act_ad": "{counter}.3 {type} {name}-3",
             }
-            .get(document_template_id.document_type_name)
+            .get(document_template_id.document_type_name, "Unknown")
             .format(
                 counter=self.counter,
                 type=_(
diff --git a/wizard/res_partner_contract_wizard.py b/wizard/res_partner_contract_wizard.py
index f8c9df1..544ea0f 100644
--- a/wizard/res_partner_contract_wizard.py
+++ b/wizard/res_partner_contract_wizard.py
@@ -226,8 +226,8 @@ class ContractWizard(models.TransientModel, Extension):
 
         self.env["mail.message"].create(
             {
-                "model": "res.partner.contract",
-                "res_id": res_id,
+                "model": self.env.context.get("attachment_model", "res.partner.contract"),
+                "res_id": self.env.context.get("attachment_res_id", res_id),
                 "message_type": "comment",
                 "attachment_ids": [(4, result.id, False)],
             }