From ec5b9a923289ef3882d89150d9eb948b7470fd16 Mon Sep 17 00:00:00 2001
From: Paul Fidika <paul@fidika.com>
Date: Tue, 14 Jul 2026 02:13:50 -0600
Subject: [PATCH] fix: preserve provider pins in slot defaults

---
 pyproject.toml                       |  2 +-
 src/gen_worker/discovery/discover.py |  8 ++++--
 tests/test_slot_discovery.py         | 41 ++++++++++++++++++++++++++++
 uv.lock                              |  2 +-
 4 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index ec33234..710a97a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "gen-worker"
-version = "0.25.0"
+version = "0.25.1"
 description = "A library used to build custom functions in Cozy Creator's serverless function platform."
 readme = "README.md"
 license = "MIT"
diff --git a/src/gen_worker/discovery/discover.py b/src/gen_worker/discovery/discover.py
index 45c7edb..0673702 100644
--- a/src/gen_worker/discovery/discover.py
+++ b/src/gen_worker/discovery/discover.py
@@ -332,8 +332,8 @@ def _stamp_family(binding_manifest: Dict[str, Any], family: str) -> None:
 
 def _model_ref_to_manifest(ref: Any) -> Dict[str, Any]:
     """``default_checkpoint``/curated-choice ref shape shared by the slots
-    block: ``{source, path, tag?, flavor?, components?}`` — a structured
-    ModelRef (pgw#511; ``components`` added pgw#505)."""
+    block: ``{source, path, tag?, flavor?, revision?, version?, components?}``
+    — a structured ModelRef (pgw#511; ``components`` added pgw#505)."""
     out: Dict[str, Any] = {"source": ref.source, "path": ref.path}
     if ref.tag and ref.tag != "latest":
         out["tag"] = ref.tag
@@ -341,6 +341,10 @@ def _model_ref_to_manifest(ref: Any) -> Dict[str, Any]:
         out["flavor"] = ref.flavor
     if ref.components:
         out["components"] = list(ref.components)
+    if ref.source in ("huggingface", "modelscope") and ref.revision:
+        out["revision"] = ref.revision
+    if ref.source == "civitai" and ref.version:
+        out["version"] = ref.version
     return out
 
 
diff --git a/tests/test_slot_discovery.py b/tests/test_slot_discovery.py
index 75ad13d..eb58b2c 100644
--- a/tests/test_slot_discovery.py
+++ b/tests/test_slot_discovery.py
@@ -74,6 +74,47 @@ def generate(self, ctx: RequestContext, data: In_) -> Out_:
     assert fn["bindings"]["vae"]["provider"] == "huggingface"
 
 
+def test_slot_preserves_provider_pins_in_binding_and_default(tmp_pkg: Path) -> None:
+    from gen_worker.discovery.discover import discover_functions
+
+    pkg = tmp_pkg / "ep_slot_pins"
+    pkg.mkdir()
+    (pkg / "__init__.py").write_text("")
+    (pkg / "main.py").write_text(textwrap.dedent("""
+        import msgspec
+        from gen_worker import Civitai, HF, ModelScope, RequestContext, Slot, endpoint
+
+        class In_(msgspec.Struct):
+            prompt: str = ""
+
+        class Out_(msgspec.Struct):
+            y: str
+
+        @endpoint(models={
+            "pipeline": Slot(object, default_checkpoint=Civitai("827184", version="2883731")),
+            "hf": Slot(object, default_checkpoint=HF("owner/repo", revision="deadbeef")),
+            "modelscope": Slot(object, default_checkpoint=ModelScope("owner/repo", revision="v1")),
+        })
+        class Gen:
+            def setup(self, pipeline: object, hf: object, modelscope: object) -> None: ...
+            def generate(self, ctx: RequestContext, data: In_) -> Out_:
+                return Out_(y="ok")
+    """))
+
+    (fn,) = discover_functions(tmp_pkg, main_module="ep_slot_pins.main")
+    slots = {slot["name"]: slot for slot in fn["slots"]}
+
+    assert fn["bindings"]["pipeline"]["ref"] == "827184"
+    assert fn["bindings"]["pipeline"]["version"] == "2883731"
+    assert slots["pipeline"]["default_checkpoint"] == {
+        "source": "civitai", "path": "827184", "version": "2883731",
+    }
+    assert fn["bindings"]["hf"]["revision"] == "deadbeef"
+    assert slots["hf"]["default_checkpoint"]["revision"] == "deadbeef"
+    assert fn["bindings"]["modelscope"]["revision"] == "v1"
+    assert slots["modelscope"]["default_checkpoint"]["revision"] == "v1"
+
+
 def test_slot_with_no_selected_by_or_default_emits_minimal_block(tmp_pkg: Path) -> None:
     from gen_worker.discovery.discover import discover_functions
 
diff --git a/uv.lock b/uv.lock
index 557378f..facd7a1 100644
--- a/uv.lock
+++ b/uv.lock
@@ -583,7 +583,7 @@ wheels = [
 
 [[package]]
 name = "gen-worker"
-version = "0.25.0"
+version = "0.25.1"
 source = { editable = "." }
 dependencies = [
     { name = "blake3" },