diff --git a/databricks/sdk/runtime/__init__.py b/databricks/sdk/runtime/__init__.py
index b79b9b2c..609b24f9 100644
--- a/databricks/sdk/runtime/__init__.py
+++ b/databricks/sdk/runtime/__init__.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 is_local_implementation = True
 
 # All objects that are injected into the Notebook's user namespace should also be made
@@ -31,15 +33,22 @@
         _globals[var] = userNamespaceGlobals[var]
     is_local_implementation = False
 except ImportError:
+    from typing import Type, cast
+
     # OSS implementation
     is_local_implementation = True
 
     try:
+        from . import stub
         from .stub import *
+        dbutils_type = Type[stub.dbutils]
     except (ImportError, NameError):
         from databricks.sdk.dbutils import RemoteDbUtils
 
         # this assumes that all environment variables are set
         dbutils = RemoteDbUtils()
+        dbutils_type = RemoteDbUtils
+
+    dbutils = cast(dbutils_type, dbutils)
 
 __all__ = ['dbutils'] if is_local_implementation else dbruntime_objects