diff --git a/strategicwm/_src/client_lib.py b/strategicwm/_src/client_lib.py
index c968231..65dd87e 100644
--- a/strategicwm/_src/client_lib.py
+++ b/strategicwm/_src/client_lib.py
@@ -86,8 +86,15 @@ def __str__(self) -> str:
             f"{self.original_http_error.args[0]}")
 
 
+RETRIABLE_EXCEPTIONS = (
+    HttpErrorRetriable,
+    genai.errors.ServerError,
+    genai.errors.APIError,
+)
+
+
 @retry.retry(
-    exceptions=HttpErrorRetriable,
+    exceptions=RETRIABLE_EXCEPTIONS,
     tries=10,
     delay=10,
     max_delay=60,
diff --git a/strategicwm/_src/se/visualization/game_tree.py b/strategicwm/_src/se/visualization/game_tree.py
index 36506a2..92d0e43 100644
--- a/strategicwm/_src/se/visualization/game_tree.py
+++ b/strategicwm/_src/se/visualization/game_tree.py
@@ -17,7 +17,10 @@
 import re
 
 import bs4
-from google.colab import output
+try:
+  from google.colab import output
+except ImportError:
+  output = None
 from IPython import display as display_html
 import ipywidgets as widgets
 import networkx as nx
@@ -138,8 +141,9 @@ def html_from_nx(self) -> str:
     soup = bs4.BeautifulSoup(html_content, "html.parser")
 
     # Register BOTH functions with unique names
-    output.register_callback("set_node_selection", self.set_node_selection)
-    output.register_callback("set_path_selection", self.set_path_selection)
+    if output:
+      output.register_callback("set_node_selection", self.set_node_selection)
+      output.register_callback("set_path_selection", self.set_path_selection)
 
     js_modifications = """
     var node_select_ts = null;
@@ -457,12 +461,13 @@ def display_html(self, html_content: str):
     html_widget = widgets.HTML(value=html_content)
     display_html.display(html_widget)
 
-    iframe_height = 2000
-    iframe_height_js = (
-        "google.colab.output.setIframeHeight(0, true, "
-        + f"{{maxHeight: {iframe_height}}})"
-    )
-    display_html.display(display_html.Javascript(iframe_height_js))
+    if output:
+      iframe_height = 2000
+      iframe_height_js = (
+          "google.colab.output.setIframeHeight(0, true, "
+          + f"{{maxHeight: {iframe_height}}})"
+      )
+      display_html.display(display_html.Javascript(iframe_height_js))
 
   def show_tree(self):
     html_content = self.html_from_nx()