Merge pull request #4 from sena-labs/claude/magical-mccarthy-9fed0e

fix: resolve 5 bugs preventing provider icons from appearing in Open WebUI

Author: sena-labs

CHANGELOG.md

@@ -10,10 +10,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - **Automatic provider-icon sync** — new `_sync_model_icons()` method writes provider icons directly into Open WebUI's Models database so they appear in the UI; controlled by the `SYNC_PROVIDER_ICONS` valve (default: enabled). Models with a manually-set icon are never overwritten
+- **`_is_owui_managed_icon()` helper** — distinguishes OWUI-default icons (`data:` URLs) and our own provider icons from user-set custom icons, enabling safe icon updates without clobbering user customisations
 
 ### Fixed
 
 - **Icon sync: correct prefixed model IDs** — `_sync_model_icons()` now discovers the pipe's `function_id` via `type(self).__module__` and writes DB records with the full prefixed ID (e.g. `openrouter_pipe.openai/gpt-4o`) matching what Open WebUI's frontend requests at `/models/model/profile/image`
+- **Icon sync: icons now actually appear in the UI** — five bugs prevented provider icons from ever showing after the first pipe load:
+  - *Wrong skip condition* — `if existing_icon:` skipped any model with *any* icon (including the generic `data:` SVG that OWUI assigns by default), so provider icons were never applied; fixed to skip only user-set custom URLs
+  - *Race condition* — `_sync_model_icons()` was called before `pipes()` returned, i.e. before OWUI registered the models; OWUI then overwrote the early insert with its own default icon; fixed by also calling `_sync_model_icons()` on cache-hit paths (until all models are confirmed synced)
+  - *Exception swallowed retry* — DB errors added the model to `_icons_synced` anyway, permanently preventing retry; removed the erroneous add
+  - *Insert marked as synced prematurely* — after `insert_new_model` the model was marked synced even though OWUI could overwrite it; the insert path no longer updates `_icons_synced`
+  - *User params clobbered* — `update_model_by_id` used an empty `ModelParams()`, erasing user-configured temperature/system-prompt/etc.; now preserves `existing.params`
+- **Icon sync: `function_id` cached at init** — `type(self).__module__` is evaluated once in `__init__` instead of on every `_sync_model_icons()` call
 - **Streaming status event** — the "done" status event is now correctly emitted at the end of streaming responses (async generator wrapper replaces sync generator that could not `await`)
 - **Dead provider-icon code removed** — `info.meta.profile_image_url` was included in model dicts returned by `pipes()` but Open WebUI ignores all fields except `id` and `name`; the field has been removed in favour of the new DB-sync approach
 - **`pipes()` response always closed** — added `finally: response.close()` to guarantee HTTP connections are returned to the session pool in all code paths (auth errors, JSON decode failures, unexpected exceptions)

openrouter_pipe.py

@@ -72,6 +72,21 @@ def _is_safe_url(url: str) -> bool:
     return isinstance(url, str) and url.lower().startswith(("http://", "https://"))
 
 
+def _is_owui_managed_icon(url: str) -> bool:
+    """Return True if the icon URL was set by OWUI or our sync logic.
+
+    data: URLs are the pipe's own SVG icon that OWUI assigns as default to all
+    manifold child models.  openrouter.ai/images/models/ URLs are the provider
+    icons we write.  Any other URL is assumed to be a user-set custom icon and
+    must not be overwritten.
+    """
+    return (
+        not url
+        or url.startswith("data:")
+        or url.startswith("https://openrouter.ai/images/models/")
+    )
+
+
 def _insert_citations(text: str, citations: Optional[List[str]]) -> str:
     """Replace [n] references with markdown links (only safe HTTP URLs)."""
     if not citations or not text:
@@ -240,6 +255,11 @@ def __init__(self) -> None:
         self._models_cache_key: str = ""
         # Track which model IDs already have icons synced (avoids repeated DB writes)
         self._icons_synced: set = set()
+        # Cache function_id once: OWUI sets __module__ to "function_{id}" at load time
+        _fm = type(self).__module__ or ""
+        self._function_id: Optional[str] = (
+            _fm[len("function_"):] if _fm.startswith("function_") else None
+        )
         if not self.valves.OPENROUTER_API_KEY:
             print("[OpenRouter Pipe] Warning: OPENROUTER_API_KEY not set")
 
@@ -274,6 +294,11 @@ def pipes(self) -> List[dict]:
 
         # Return cached models if still valid
         if self._models_cache_valid() and self._models_cache is not None:
+            # Continue syncing icons on cache hits until all models are confirmed.
+            # This resolves the race condition where OWUI registers models (and may
+            # overwrite icons) only after the first pipes() call returns.
+            if self.valves.SYNC_PROVIDER_ICONS and len(self._icons_synced) < len(self._models_cache):
+                self._sync_model_icons(self._models_cache)
             return self._models_cache
 
         headers = self._build_headers(include_content_type=False)
@@ -456,14 +481,18 @@ def _sync_model_icons(self, models: List[dict]) -> None:
         """Write provider icons into Open WebUI's Models DB.
 
         Open WebUI serves model icons from its database, not from the dicts
-        returned by ``pipes()``.  Crucially, Open WebUI prefixes every pipe
-        model ID with ``{function_id}.`` (e.g. ``openrouter_pipe.openai/gpt-4o``)
-        and the frontend requests icons using that prefixed ID.  This method
-        discovers the pipe's own *function_id* from ``type(self).__module__``
-        (set to ``function_{id}`` by Open WebUI's module loader) and writes
-        DB records keyed by the full prefixed ID.
-
-        Models that already have a custom icon in the DB are skipped.
+        returned by ``pipes()``.  OWUI prefixes every pipe model ID with
+        ``{function_id}.`` (e.g. ``openrouter_pipe.openai/gpt-4o``) and the
+        frontend requests icons using that prefixed ID.
+
+        Called both on cache miss and on subsequent cache hits (until all
+        models are confirmed synced).  The cache-hit path is needed because
+        OWUI registers models *after* ``pipes()`` returns, potentially
+        overwriting any early insert with its own default icon; the second
+        call finds the models already in DB and updates them correctly.
+
+        User-set custom icons (any URL that is not a ``data:`` URL and does not
+        start with ``https://openrouter.ai/images/models/``) are preserved.
         This is a best-effort operation — failures are silently logged.
         """
         try:
@@ -477,14 +506,10 @@ def _sync_model_icons(self, models: List[dict]) -> None:
             # Running outside Open WebUI (e.g. standalone tests) — skip silently
             return
 
-        # Discover the pipe's function_id.  Open WebUI loads pipe modules as
-        # ``function_{function_id}`` so type(self).__module__ exposes it.
-        func_module = type(self).__module__ or ""
-        if func_module.startswith("function_"):
-            function_id = func_module[len("function_"):]
-        else:
-            # Cannot determine the pipe's function_id — skip icon sync
+        # function_id was resolved once in __init__ from type(self).__module__
+        if not self._function_id:
             return
+        function_id = self._function_id
 
         for model in models:
             model_id = model.get("id", "")
@@ -509,17 +534,28 @@ def _sync_model_icons(self, models: List[dict]) -> None:
             try:
                 existing = Models.get_model_by_id(db_model_id)
                 if existing:
-                    # If model already has a custom icon, don't overwrite it
                     existing_icon = ""
                     if hasattr(existing, "meta") and existing.meta:
                         existing_icon = (
                             getattr(existing.meta, "profile_image_url", "") or ""
                         )
-                    if existing_icon:
+
+                    # Skip if icon is already the correct provider URL
+                    if existing_icon == icon_url:
+                        self._icons_synced.add(model_id)
+                        continue
+
+                    # Skip if icon was set by the user (not by OWUI or our sync).
+                    # data: URLs are OWUI defaults; openrouter.ai URLs are ours.
+                    if existing_icon and not _is_owui_managed_icon(existing_icon):
                         self._icons_synced.add(model_id)
                         continue
 
-                    # Update existing model with icon
+                    # Proceed: icon is empty, an OWUI default, or one of our URLs
+                    # Update existing model with icon, preserving user-set params
+                    existing_params = ModelParams()
+                    if hasattr(existing, "params") and existing.params:
+                        existing_params = existing.params
                     Models.update_model_by_id(
                         db_model_id,
                         ModelForm(
@@ -530,26 +566,34 @@ def _sync_model_icons(self, models: List[dict]) -> None:
                                 else model.get("name", model_id)
                             ),
                             meta=ModelMeta(profile_image_url=icon_url),
-                            params=ModelParams(),
+                            params=existing_params,
                         ),
                     )
                 else:
-                    # Insert new model record with icon
-                    Models.insert_new_model(
-                        ModelForm(
-                            id=db_model_id,
-                            name=model.get("name", model_id),
-                            meta=ModelMeta(profile_image_url=icon_url),
-                            params=ModelParams(),
-                        ),
-                        user_id="pipe:openrouter",
-                    )
+                    # Model not yet in DB — best-effort early insert.
+                    # OWUI will register models after pipes() returns and may
+                    # overwrite this record, so do NOT mark as synced here.
+                    # The next cache-hit call to _sync_model_icons will find the
+                    # model in DB and update it correctly.
+                    try:
+                        Models.insert_new_model(
+                            ModelForm(
+                                id=db_model_id,
+                                name=model.get("name", model_id),
+                                meta=ModelMeta(profile_image_url=icon_url),
+                                params=ModelParams(),
+                            ),
+                            user_id="pipe:openrouter",
+                        )
+                    except Exception:
+                        pass
+                    continue  # do not add to _icons_synced yet
 
                 self._icons_synced.add(model_id)
             except Exception as exc:
                 # Best-effort — don't let icon sync break model listing
+                # Do NOT add to _icons_synced: allow retry on next call
                 print(f"[OpenRouter Pipe] Icon sync failed for {db_model_id}: {exc}")
-                self._icons_synced.add(model_id)  # Don't retry on every refresh
 
     @staticmethod
     def get_provider_icon(provider: str) -> Optional[str]:

test_pipe.py

@@ -43,6 +43,7 @@
 _insert_citations = mod._insert_citations
 _format_citation_list = mod._format_citation_list
 _OWUI_INTERNAL_KEYS = mod._OWUI_INTERNAL_KEYS
+_is_owui_managed_icon = mod._is_owui_managed_icon
 
 # ── Helpers ───────────────────────────────────────────────────────────────────
 _PASS = 0
@@ -1380,7 +1381,9 @@ async def _test_pipe_no_msgs_key():
 # 25d. Valve default is True
 _assert(Pipe.Valves(OPENROUTER_API_KEY="k").SYNC_PROVIDER_ICONS is True, "SYNC_PROVIDER_ICONS default is True")
 
-# 25e. Non-function module → returns early without DB calls
+# 25e. No function_id → returns early without DB calls
+# function_id is cached in __init__; pipes created in tests have _function_id=None
+# because the test module name doesn't start with "function_".
 _pipe_nofunc = Pipe()
 _pipe_nofunc.valves = Pipe.Valves(OPENROUTER_API_KEY="k", SYNC_PROVIDER_ICONS=True)
 _mock_Models_nf = MagicMock()
@@ -1389,25 +1392,25 @@ async def _test_pipe_no_msgs_key():
 _fake_owui_nf.ModelForm = MagicMock()
 _fake_owui_nf.ModelMeta = MagicMock()
 _fake_owui_nf.ModelParams = MagicMock()
-_orig_module_nf = Pipe.__module__
 try:
     sys.modules["open_webui.models.models"] = _fake_owui_nf
-    # Module name doesn't start with "function_" → should skip
-    Pipe.__module__ = "openrouter_pipe"
+    # _pipe_nofunc._function_id is None (no "function_" prefix at init time) → skip
+    _assert(_pipe_nofunc._function_id is None, "_sync_model_icons: _function_id is None outside OWUI")
     _pipe_nofunc._sync_model_icons([{"id": "openai/gpt-4o", "name": "GPT-4o"}])
     _assert(
         not _mock_Models_nf.get_model_by_id.called,
-        "_sync_model_icons: skips DB when module is not function_*",
+        "_sync_model_icons: skips DB when _function_id is None",
     )
 finally:
-    Pipe.__module__ = _orig_module_nf
     sys.modules.pop("open_webui.models.models", None)
 
-# 25f. With function_ module → uses prefixed IDs for DB operations
+# 25f. With function_id set → uses prefixed IDs; does NOT add to _icons_synced after insert
+# (OWUI may overwrite the record after pipes() returns; confirmed on next call)
 _pipe_func = Pipe()
 _pipe_func.valves = Pipe.Valves(OPENROUTER_API_KEY="k", SYNC_PROVIDER_ICONS=True)
+_pipe_func._function_id = "openrouter_pipe"  # simulate OWUI module naming
 _mock_Models_f = MagicMock()
-_mock_Models_f.get_model_by_id.return_value = None  # No existing record
+_mock_Models_f.get_model_by_id.return_value = None  # No existing record yet
 _mock_ModelForm_f = MagicMock()
 _mock_ModelMeta_f = MagicMock()
 _mock_ModelParams_f = MagicMock()
@@ -1416,10 +1419,8 @@ async def _test_pipe_no_msgs_key():
 _fake_owui_f.ModelForm = _mock_ModelForm_f
 _fake_owui_f.ModelMeta = _mock_ModelMeta_f
 _fake_owui_f.ModelParams = _mock_ModelParams_f
-_orig_module_f = Pipe.__module__
 try:
     sys.modules["open_webui.models.models"] = _fake_owui_f
-    Pipe.__module__ = "function_openrouter_pipe"
     _pipe_func._sync_model_icons([
         {"id": "openai/gpt-4o", "name": "GPT-4o"},
         {"id": "anthropic/claude-3.5-sonnet", "name": "Claude 3.5"},
@@ -1440,9 +1441,6 @@ async def _test_pipe_no_msgs_key():
         "_sync_model_icons: insert_new_model called for new models",
     )
     # Verify the ModelForm ID is prefixed
-    _insert_form = _mock_Models_f.insert_new_model.call_args_list[0].args[0]
-    _insert_id = _mock_ModelForm_f.call_args_list[0].kwargs.get("id", _mock_ModelForm_f.call_args_list[0].args[0] if _mock_ModelForm_f.call_args_list[0].args else "")
-    # ModelForm was called with id=prefixed_id
     _form_calls = _mock_ModelForm_f.call_args_list
     _form_ids = [c.kwargs.get("id", "") for c in _form_calls]
     _assert(
@@ -1453,45 +1451,89 @@ async def _test_pipe_no_msgs_key():
         "openrouter_pipe.anthropic/claude-3.5-sonnet" in _form_ids,
         "_sync_model_icons: ModelForm uses prefixed ID (anthropic)",
     )
-    # Verify _icons_synced tracks the RAW model IDs
+    # After insert (model not yet registered by OWUI), _icons_synced must NOT be updated.
+    # The next cache-hit call will confirm the icon is set correctly.
     _assert(
-        "openai/gpt-4o" in _pipe_func._icons_synced,
-        "_sync_model_icons: _icons_synced uses raw model ID",
+        "openai/gpt-4o" not in _pipe_func._icons_synced,
+        "_sync_model_icons: _icons_synced NOT updated after insert (allows retry)",
     )
 finally:
-    Pipe.__module__ = _orig_module_f
     sys.modules.pop("open_webui.models.models", None)
 
-# 25g. Existing model with icon → skips overwrite
+# 25g. Existing model with user custom icon → skips overwrite
 _pipe_skip = Pipe()
 _pipe_skip.valves = Pipe.Valves(OPENROUTER_API_KEY="k", SYNC_PROVIDER_ICONS=True)
+_pipe_skip._function_id = "openrouter_pipe"  # simulate OWUI module naming
 _mock_Models_s = MagicMock()
 _existing_model = MagicMock()
 _existing_model.meta.profile_image_url = "https://custom-icon.example.com/icon.png"
 _existing_model.name = "Custom GPT"
+_existing_model.params = None
 _mock_Models_s.get_model_by_id.return_value = _existing_model
 _fake_owui_s = ModuleType("open_webui.models.models")
 _fake_owui_s.Models = _mock_Models_s
 _fake_owui_s.ModelForm = MagicMock()
 _fake_owui_s.ModelMeta = MagicMock()
 _fake_owui_s.ModelParams = MagicMock()
-_orig_module_s = Pipe.__module__
 try:
     sys.modules["open_webui.models.models"] = _fake_owui_s
-    Pipe.__module__ = "function_openrouter_pipe"
     _pipe_skip._sync_model_icons([{"id": "openai/gpt-4o", "name": "GPT-4o"}])
     _assert(
         not _mock_Models_s.update_model_by_id.called,
-        "_sync_model_icons: skips update when model has custom icon",
+        "_sync_model_icons: skips update when model has user custom icon",
     )
     _assert(
         not _mock_Models_s.insert_new_model.called,
-        "_sync_model_icons: skips insert when model has custom icon",
+        "_sync_model_icons: skips insert when model has user custom icon",
+    )
+    _assert(
+        "openai/gpt-4o" in _pipe_skip._icons_synced,
+        "_sync_model_icons: adds to _icons_synced when skipping custom icon",
     )
 finally:
-    Pipe.__module__ = _orig_module_s
     sys.modules.pop("open_webui.models.models", None)
 
+# 25h. Existing model with OWUI default (data: URL) icon → updates with provider icon
+_pipe_update = Pipe()
+_pipe_update.valves = Pipe.Valves(OPENROUTER_API_KEY="k", SYNC_PROVIDER_ICONS=True)
+_pipe_update._function_id = "openrouter_pipe"  # simulate OWUI module naming
+_mock_Models_u = MagicMock()
+_existing_default = MagicMock()
+_existing_default.name = "GPT-4o"
+_existing_default.meta.profile_image_url = "data:image/svg+xml;base64,ABC123=="
+_existing_default.params = None
+_mock_Models_u.get_model_by_id.return_value = _existing_default
+_fake_owui_u = ModuleType("open_webui.models.models")
+_fake_owui_u.Models = _mock_Models_u
+_fake_owui_u.ModelForm = MagicMock()
+_fake_owui_u.ModelMeta = MagicMock()
+_fake_owui_u.ModelParams = MagicMock()
+try:
+    sys.modules["open_webui.models.models"] = _fake_owui_u
+    _pipe_update._sync_model_icons([{"id": "openai/gpt-4o", "name": "GPT-4o"}])
+    _assert(
+        _mock_Models_u.update_model_by_id.called,
+        "_sync_model_icons: updates model when icon is OWUI default (data: URL)",
+    )
+    _assert(
+        not _mock_Models_u.insert_new_model.called,
+        "_sync_model_icons: does not insert when model already exists",
+    )
+    _assert(
+        "openai/gpt-4o" in _pipe_update._icons_synced,
+        "_sync_model_icons: adds to _icons_synced after successful update",
+    )
+finally:
+    sys.modules.pop("open_webui.models.models", None)
+
+# 25i. _is_owui_managed_icon helper
+_is_owui = mod._is_owui_managed_icon
+_assert(_is_owui(""), "_is_owui_managed_icon: empty string → True (no icon)")
+_assert(_is_owui("data:image/svg+xml;base64,ABC"), "_is_owui_managed_icon: data: URL → True")
+_assert(_is_owui("https://openrouter.ai/images/models/openai.svg"), "_is_owui_managed_icon: openrouter.ai URL → True")
+_assert(not _is_owui("https://custom-icon.example.com/icon.png"), "_is_owui_managed_icon: external URL → False")
+_assert(not _is_owui("https://cdn.openai.com/logo.png"), "_is_owui_managed_icon: other https URL → False")
+
 # ══════════════════════════════════════════════════════════════════════════════
 # Summary
 # ══════════════════════════════════════════════════════════════════════════════