diff --git a/Language/Haskell/GhcMod/Gap.hs b/Language/Haskell/GhcMod/Gap.hs
index a2024aa..31555a8 100644
--- a/Language/Haskell/GhcMod/Gap.hs
+++ b/Language/Haskell/GhcMod/Gap.hs
@@ -8,7 +8,7 @@ module Language.Haskell.GhcMod.Gap (
   , supportedExtensions
   , getSrcSpan
   , getSrcFile
-  , setCtx
+  , withContext
   , fOptions
   , toStringBuffer
   , showSeverityCaption
@@ -32,6 +32,7 @@ module Language.Haskell.GhcMod.Gap (
   , showDocWith
   , GapThing(..)
   , fromTyThing
+  , fileModSummary
   ) where
 
 import Control.Applicative hiding (empty)
@@ -187,34 +188,38 @@ fOptions = [option | (option,_,_) <- fFlags]
 ----------------------------------------------------------------
 ----------------------------------------------------------------
 
-setCtx :: FilePath -> Ghc ModSummary
-#if __GLASGOW_HASKELL__ >= 704
-setCtx file = do
+fileModSummary :: FilePath -> Ghc ModSummary
+fileModSummary file = do
+    mss <- getModuleGraph
+    let [ms] = filter (\m -> ml_hs_file (ms_location m) == Just file) mss
+    return ms
+
+withContext :: Ghc a -> Ghc a
+withContext action = gbracket setup teardown body
+  where
+    setup = getContext
+    teardown = setContext
+    body _ = do
+        topImports >>= setContext
+        action
+
+topImports :: Ghc [InteractiveImport]
+topImports = do
     mss <- getModuleGraph
 #if __GLASGOW_HASKELL__ >= 706
     let modName = IIModule . moduleName . ms_mod
-#else
+#elif __GLASGOW_HASKELL__ >= 704
     let modName = IIModule . ms_mod
-#endif
-    top <- map modName <$> filterM isTop mss
-    setContext top
-    let [ms] = filter (\m -> ml_hs_file (ms_location m) == Just file) mss
-    return ms
 #else
-setCtx file = do
-    mss <- getModuleGraph
-    top <- map ms_mod <$> filterM isTop mss
-    setContext top []
-    let [ms] = filter (\m -> ml_hs_file (ms_location m) == Just file) mss
-    return ms
+    let modName = ms_mod
 #endif
+    map modName <$> filterM isTop mss
   where
     isTop mos = lookupMod ||> returnFalse
       where
         lookupMod = lookupModule (ms_mod_name mos) Nothing >> return True
         returnFalse = return False
 
-
 showSeverityCaption :: Severity -> String
 #if __GLASGOW_HASKELL__ >= 706
 showSeverityCaption SevWarning = "Warning: "
diff --git a/Language/Haskell/GhcMod/Info.hs b/Language/Haskell/GhcMod/Info.hs
index 0c61b2a..416ec23 100644
--- a/Language/Haskell/GhcMod/Info.hs
+++ b/Language/Haskell/GhcMod/Info.hs
@@ -47,7 +47,7 @@ info :: Options
      -> Ghc String
 info opt file expr = convert opt <$> ghandle handler body
   where
-    body = inModuleContext file $ \_ dflag style -> do
+    body = inModuleContext file $ \dflag style -> do
         sdoc <- Gap.infoThing expr
         return $ showPage dflag style sdoc
     handler (SomeException _) = return "Cannot show info"
@@ -84,7 +84,8 @@ types :: Options
       -> Ghc String
 types opt file lineNo colNo = convert opt <$> ghandle handler body
   where
-    body = inModuleContext file $ \modSum dflag style -> do
+    body = inModuleContext file $ \dflag style -> do
+        modSum <- Gap.fileModSummary file
         srcSpanTypes <- getSrcSpanType modSum lineNo colNo
         return $ map (toTup dflag style) $ sortBy (cmp `on` fst) srcSpanTypes
     handler (SomeException _) = return []
@@ -126,10 +127,10 @@ pretty dflag style = showOneLine dflag style . Gap.typeForUser
 
 ----------------------------------------------------------------
 
-inModuleContext :: FilePath -> (G.ModSummary -> DynFlags -> PprStyle -> Ghc a) -> Ghc a
+inModuleContext :: FilePath -> (DynFlags -> PprStyle -> Ghc a) -> Ghc a
 inModuleContext file action = withDynFlags setDeferTypeErrors $ do
     setTargetFiles [file]
-    modSum <- Gap.setCtx file
-    dflag <- G.getSessionDynFlags
-    style <- getStyle
-    action modSum dflag style
+    Gap.withContext $ do
+        dflag <- G.getSessionDynFlags
+        style <- getStyle
+        action dflag style