diff --git a/src/lib_context/context.mli b/src/lib_context/context.mli
index 169cc6812be0b1b041ed8c91bdd8477ce22b1145..9109d21d056a2187e911082c60bd3c91385cd5dc 100644
--- a/src/lib_context/context.mli
+++ b/src/lib_context/context.mli
@@ -72,6 +72,8 @@ module Make (Encoding : module type of Tezos_context_encoding.Context) : sig
     Does not fail when the context is not in read-only mode. *)
   val sync : index -> unit Lwt.t
 
+  val flush : t -> t Lwt.t
+
   val compute_testchain_chain_id : Block_hash.t -> Chain_id.t
 
   val compute_testchain_genesis : Block_hash.t -> Block_hash.t
diff --git a/src/lib_context/helpers/context.ml b/src/lib_context/helpers/context.ml
index 73e1b3cbc45561cd0cbc9c238157050b4d426bd4..634be779dc686e1fc8bbcc67050f74ab95fe1a0b 100644
--- a/src/lib_context/helpers/context.ml
+++ b/src/lib_context/helpers/context.ml
@@ -199,6 +199,11 @@ module Make_tree (Conf : Conf) (Store : DB) = struct
     in
     Store.Tree.shallow repo kinded_hash
 
+  let kinded_key t =
+    match Store.Tree.key t with
+    | (None | Some (`Node _)) as r -> r
+    | Some (`Contents (v, ())) -> Some (`Value v)
+
   let is_shallow tree =
     match Store.Tree.inspect tree with
     | `Node `Key -> true
diff --git a/src/lib_context/helpers/context.mli b/src/lib_context/helpers/context.mli
index eecf974084415e2bdd47b0ca284c84a616f509e5..3b8c99151fe9837b6da6adcdbbd9bd2ecb7e2c83 100644
--- a/src/lib_context/helpers/context.mli
+++ b/src/lib_context/helpers/context.mli
@@ -58,6 +58,8 @@ module Make_tree (Conf : Conf) (DB : DB) : sig
 
   val shallow : DB.repo -> kinded_key -> DB.tree
 
+  val kinded_key : DB.tree -> kinded_key option
+
   val is_shallow : DB.tree -> bool
 
   (** Exception raised by [find_tree] and [add_tree] when applied to shallow
diff --git a/src/lib_context/memory/context.ml b/src/lib_context/memory/context.ml
index f60b8176110f7a02fbc76b46358f222439832325..7eea8599552d9066e66fa7a494ccc23884bfdd10 100644
--- a/src/lib_context/memory/context.ml
+++ b/src/lib_context/memory/context.ml
@@ -40,12 +40,22 @@ module Make (Encoding : module type of Tezos_context_encoding.Context) = struct
       match t with
       | `Node hash -> `Node (Hash.of_context_hash hash)
       | `Value hash -> `Contents (Hash.of_context_hash hash, ())
+
+    let of_irmin_key t : kinded_key =
+      match t with
+      | `Node hash -> `Node (Hash.to_context_hash hash)
+      | `Contents (hash, ()) -> `Value (Hash.to_context_hash hash)
   end
 
   module Tree = struct
     include Tezos_context_helpers.Context.Make_tree (Conf) (Store)
 
     let shallow repo key = Store.Tree.shallow repo (Kinded_key.to_irmin_key key)
+
+    let kinded_key tree =
+      match Store.Tree.key tree with
+      | None -> None
+      | Some h -> Some (Kinded_key.of_irmin_key h)
   end
 
   include Tree
diff --git a/src/lib_context/sigs/context.ml b/src/lib_context/sigs/context.ml
index 341c915a8e9dd987bdd8ab06affd039b618ef931..7a636de2698da5a99ebaa6404718ab6ede93b528 100644
--- a/src/lib_context/sigs/context.ml
+++ b/src/lib_context/sigs/context.ml
@@ -453,6 +453,8 @@ module type S = sig
     val shallow : repo -> kinded_key -> tree
 
     val is_shallow : tree -> bool
+
+    val kinded_key : tree -> kinded_key option
   end
 
   (** [produce r h f] runs [f] on top of a real store [r], producing a proof and