diff --git a/.travis.yml b/.travis.yml
index 486451106f..b53888f649 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -12,7 +12,7 @@ python:
     - "3.6"
 
 install:
-    - pip install demjson
+    - pip install demjson graphviz
     - pip install -r tools/requirements.txt
 
 script:
@@ -20,6 +20,7 @@ script:
     - jsonlint defs/*/*.json
     - python tools/cointool.py check
     - python tools/support.py check --ignore-tokens --ignore-missing
+    - python protob/graph.py protob/*.proto
 
 notifications:
   webhooks:
diff --git a/protob/graph.py b/protob/graph.py
index 6b6543e0a7..362a1d8b1a 100755
--- a/protob/graph.py
+++ b/protob/graph.py
@@ -1,11 +1,13 @@
 #!/usr/bin/env python3
 import sys
+from os.path import basename
 
 from graphviz import Digraph
 
 
 class Message(object):
-    def __init__(self, name, attrs):
+    def __init__(self, fname, name, attrs):
+        self.fname = basename(fname)
         self.name = name
         if len(attrs) == 0:
             raise ValueError("message '%s' has no attributes" % name)
@@ -16,18 +18,20 @@ class Message(object):
         elif t == "next":
             self.typ = "normal"
             attrs = attrs
-        elif t == "wrap":
-            self.typ = "normal"
-            attrs = attrs
         else:
             raise ValueError("wrong message type in message '%s'" % name)
         self.next = []
-        self.wrap = []
         for a in attrs:
             if a[0] == "next":
                 self.next.append(a[1])
-            elif a[0] == "wrap":
-                self.wrap.append(a[1])
+
+    def __repr__(self):
+        return '%s(type=%s, fname="%s", next=%s)' % (
+            self.name,
+            self.typ,
+            self.fname,
+            self.next,
+        )
 
 
 def generate_messages(files):
@@ -40,13 +44,13 @@ def generate_messages(files):
                 attrs.append(line[4:].split(" "))
             elif line.startswith("message "):
                 name = line[8:-2]
-                msgs[name] = Message(name, attrs)
+                msgs[name] = Message(f, name, attrs)
                 attrs = []
     return msgs
 
 
 def generate_graph(msgs, fn):
-    dot = Digraph(format="png")
+    dot = Digraph()
     dot.attr(rankdir="LR")
     for m in msgs.values():
         if m.typ == "start":
@@ -63,10 +67,15 @@ def generate_graph(msgs, fn):
     for m in msgs.values():
         for n in m.next:
             dot.edge(m.name, n)
-        for n in m.wrap:
-            dot.edge(m.name, n)
     dot.render(fn)
 
 
-msgs = generate_messages(sys.argv)
-generate_graph(msgs, "graph.gv")
+def main():
+    proto_files = sys.argv
+
+    msgs = generate_messages(proto_files)
+    generate_graph(msgs, "proto.gv")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/protob/messages-monero.proto b/protob/messages-monero.proto
index a900f95a0b..2f5b812c58 100644
--- a/protob/messages-monero.proto
+++ b/protob/messages-monero.proto
@@ -54,6 +54,7 @@ message MoneroTransactionDestinationEntry {
 
 /**
  * Range sig parameters / data.
+ * @embed
  */
 message MoneroTransactionRsigData {
     optional uint32 rsig_type = 1;  // range signature (aka proof) type
diff --git a/protob/messages-ontology.proto b/protob/messages-ontology.proto
index cac05cd8e0..1bb1845693 100644
--- a/protob/messages-ontology.proto
+++ b/protob/messages-ontology.proto
@@ -63,7 +63,7 @@ message OntologyAddress {
 
 /**
  * Request: Ask device to sign Ontology transfer
- * @start
+ * @start
  * @next OntologySignedTransfer
  */
 message OntologySignTransfer {
@@ -99,7 +99,7 @@ message OntologySignedTransfer {
 
 /**
  * Request: Ask device to sign Ontology ONG withdrawal
- * @start
+ * @start
  * @next OntologySignedWithdrawOng
  */
 message OntologySignWithdrawOng {
@@ -127,7 +127,7 @@ message OntologySignedWithdrawOng {
 
 /**
  * Request: Ask device to sign Ontology ONT ID registration
- * @start
+ * @start
  * @next OntologySignedOntIdRegister
  */
 message OntologySignOntIdRegister {
@@ -154,7 +154,7 @@ message OntologySignedOntIdRegister {
 
 /**
  * Request: Ask device to sign Ontology ONT ID attributes adding
- * @start
+ * @start
  * @next OntologySignedOntIdAddAttributes
  */
 message OntologySignOntIdAddAttributes {
diff --git a/protob/messages-tron.proto b/protob/messages-tron.proto
index b297eeaffc..ce3525c93c 100644
--- a/protob/messages-tron.proto
+++ b/protob/messages-tron.proto
@@ -6,7 +6,7 @@ option java_outer_classname = "TrezorMessageTron";
 
 /**
  * Request: Ask device for Tron address corresponding to address_n path
- * @next PassphraseRequest
+ * @start
  * @next TronAddress
  * @next Failure
  */
@@ -17,7 +17,7 @@ message TronGetAddress {
 
 /**
  * Response: Contains Tron address derived from device private seed
- * @prev TronGetAddress
+ * @end
  */
 message TronAddress {
     optional string address = 1;                                // Tron address (base58)
@@ -25,12 +25,13 @@ message TronAddress {
 
 /**
  * Request: Ask device to sign Tron transaction
+ * @start
  * @next TronSignedTx
  */
 message TronSignTx {
     repeated uint32 address_n = 1;                              // BIP-32 path to derive the key from master node
     // Common part of transaction
-    optional bytes ref_block_bytes = 2;                         // Reference block number 
+    optional bytes ref_block_bytes = 2;                         // Reference block number
     optional bytes ref_block_hash = 3;                          // Reference block hash
     optional uint64 expiration = 4;                             // Transaction expiration
     optional string data = 5;                                   // Extra transaction info
@@ -90,13 +91,13 @@ message TronSignTx {
             optional string description = 10;                   // Asset description
             optional string url = 11;                           // Asset URL
         }
-        // Participate in an asset 
+        // Participate in an asset
         message TronParticipateAssetIssueContract {
             optional bytes to_address = 1;                      // Asset issuer address - decoded base 58
             optional string asset_name = 2;                     // The name of target asset
             optional uint64 amount = 3;                         // TRX amount in sun
         }
-        // Freeze TRX balance                        
+        // Freeze TRX balance
         message TronFreezeBalanceContract {
             optional uint64 frozen_balance = 1;                 // Amount to freeze
             optional uint64 frozen_duration = 2;                // Freeze minimal duration in days
@@ -118,7 +119,7 @@ message TronSignTx {
         // Network proposal contract
         message TronProposalCreateContract {
             message TronProposalParameters {
-                optional uint64 key = 1;                        // Paramter key
+                optional uint64 key = 1;                        // Parameter key
                 optional uint64 value = 2;                      // Parameter value
             }
             repeated TronProposalParameters parameters = 1;     // Parameter to be changed
@@ -132,7 +133,6 @@ message TronSignTx {
         message TronProposalDeleteContract {
             optional uint64 proposal_id = 1;                    // Proposal ID
         }
-        
         optional TronTransferContract transfer_contract = 1;
         optional TronTransferAssetContract transfer_asset_contract = 2;
         optional TronVoteWitnessContract vote_witness_contract = 4;
@@ -153,9 +153,9 @@ message TronSignTx {
 }
 
 /**
-* Response: Contains Tron transaction signature
-* @prev TronSignTx
-*/
+ * Response: Contains Tron transaction signature
+ * @end
+ */
 message TronSignedTx {
     optional bytes signature = 1;                               // Transaction signature
     optional bytes serialized_tx = 2;                           // Serialized transaction