From dedfb36ff221e6f64deb2f720dccd5d94949b248 Mon Sep 17 00:00:00 2001
From: Ousama Ben Younes <benyounes.ousama@gmail.com>
Date: Tue, 7 Jul 2026 20:10:06 +0000
Subject: [PATCH] fix(providers): show vertex extra hint for wrapped import
 errors

---
 strix/interface/main.py      | 44 +++++++++++++++++++++++++++---------
 tests/test_provider_hints.py | 28 ++++++++++++++++++-----
 2 files changed, 55 insertions(+), 17 deletions(-)

diff --git a/strix/interface/main.py b/strix/interface/main.py
index 8bd65fbf..39758f31 100644
--- a/strix/interface/main.py
+++ b/strix/interface/main.py
@@ -56,6 +56,16 @@ from strix.telemetry.logging import configure_dependency_logging
 
 
 HOST_GATEWAY_HOSTNAME = "host.docker.internal"
+BEDROCK_MODEL_PREFIX = "bedrock/"
+BEDROCK_MISSING_MODULE_ERROR = "No module named 'boto3'"
+BEDROCK_EXTRA_HINT = (
+    'Bedrock support is optional. Install it with: pipx install "strix-agent[bedrock]"'
+)
+VERTEX_MODEL_MARKER = "vertex"
+VERTEX_MISSING_MODULE_ERROR = "No module named 'google'"
+VERTEX_EXTRA_HINT = (
+    'Vertex AI support is optional. Install it with: pipx install "strix-agent[vertex]"'
+)
 
 
 import logging  # noqa: E402
@@ -214,23 +224,35 @@ def check_docker_installed() -> None:
     logger.debug("Docker CLI present")
 
 
+def _exception_messages(exc: BaseException) -> tuple[str, ...]:
+    messages = [str(exc)]
+    if exc.__cause__ is not None:
+        messages.append(str(exc.__cause__))
+    if exc.__context__ is not None:
+        messages.append(str(exc.__context__))
+    return tuple(messages)
+
+
 def _provider_import_hint(exc: BaseException, model: str) -> str | None:
     """Return an install hint when *exc* is a missing provider dependency.
 
     Bedrock and Vertex AI ship as optional extras: Bedrock needs ``boto3`` and
-    Vertex AI needs ``google-auth``. When either is absent, litellm raises an
-    ``ImportError``/``ModuleNotFoundError`` naming the missing package. Map that
-    back to the matching extra so the user knows what to install. Returns
-    ``None`` for any unrelated error.
+    Vertex AI needs ``google-auth``. When either is absent, litellm may raise an
+    ``ImportError``/``ModuleNotFoundError`` directly or wrap it in a connection
+    error. Map the missing module back to the matching extra so the user knows
+    what to install. Returns ``None`` for any unrelated error.
     """
-    if not isinstance(exc, ImportError):
-        return None
-    message = str(exc)
     model_name = model.lower()
-    if "boto3" in message and model_name.startswith("bedrock/"):
-        return 'Bedrock support is optional. Install it with: pipx install "strix-agent[bedrock]"'
-    if "google" in message and "vertex" in model_name:
-        return 'Vertex AI support is optional. Install it with: pipx install "strix-agent[vertex]"'
+    messages = _exception_messages(exc)
+    if any(
+        BEDROCK_MISSING_MODULE_ERROR in message for message in messages
+    ) and model_name.startswith(BEDROCK_MODEL_PREFIX):
+        return BEDROCK_EXTRA_HINT
+    if (
+        any(VERTEX_MISSING_MODULE_ERROR in message for message in messages)
+        and VERTEX_MODEL_MARKER in model_name
+    ):
+        return VERTEX_EXTRA_HINT
     return None
 
 
diff --git a/tests/test_provider_hints.py b/tests/test_provider_hints.py
index a71ade99..4ec4366c 100644
--- a/tests/test_provider_hints.py
+++ b/tests/test_provider_hints.py
@@ -5,20 +5,36 @@ from __future__ import annotations
 from strix.interface.main import _provider_import_hint
 
 
+VERTEX_MODEL = "vertex_ai/gemini-3-pro-preview"
+BEDROCK_MODEL = "bedrock/anthropic.claude-4-5-sonnet"
+VERTEX_EXTRA_NAME = "vertex"
+BEDROCK_EXTRA_NAME = "bedrock"
+INSTALL_EXTRA_COMMAND_FRAGMENT = 'pipx install "strix-agent['
+WRAPPED_VERTEX_GOOGLE_ERROR = "litellm.APIConnectionError: No module named 'google'"
+
+
 def test_bedrock_boto3_hint() -> None:
     exc = ModuleNotFoundError("No module named 'boto3'")
-    hint = _provider_import_hint(exc, "bedrock/anthropic.claude-4-5-sonnet")
+    hint = _provider_import_hint(exc, BEDROCK_MODEL)
     assert hint is not None
-    assert 'pipx install "strix-agent[' in hint
-    assert "bedrock" in hint
+    assert INSTALL_EXTRA_COMMAND_FRAGMENT in hint
+    assert BEDROCK_EXTRA_NAME in hint
 
 
 def test_vertex_google_hint() -> None:
     exc = ImportError("No module named 'google'")
-    hint = _provider_import_hint(exc, "vertex_ai/gemini-3-pro-preview")
+    hint = _provider_import_hint(exc, VERTEX_MODEL)
     assert hint is not None
-    assert 'pipx install "strix-agent[' in hint
-    assert "vertex" in hint
+    assert INSTALL_EXTRA_COMMAND_FRAGMENT in hint
+    assert VERTEX_EXTRA_NAME in hint
+
+
+def test_vertex_google_hint_for_litellm_wrapped_connection_error() -> None:
+    exc = ConnectionError(WRAPPED_VERTEX_GOOGLE_ERROR)
+    hint = _provider_import_hint(exc, VERTEX_MODEL)
+    assert hint is not None
+    assert INSTALL_EXTRA_COMMAND_FRAGMENT in hint
+    assert VERTEX_EXTRA_NAME in hint
 
 
 def test_non_import_error_returns_none() -> None: