From 34d6a922c2c5ffcc562dfa50a54c6d5de818a1bd Mon Sep 17 00:00:00 2001
From: Owen Kephart <owenkephart@gmail.com>
Date: Fri, 10 Jan 2025 17:15:48 -0800
Subject: [PATCH] [components][fix] Fix issue with additional scope
 determination

---
 .../dagster_components/core/schema/metadata.py                | 4 +++-
 .../rendering_tests/test_schema_resolution.py                 | 4 ++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/python_modules/libraries/dagster-components/dagster_components/core/schema/metadata.py b/python_modules/libraries/dagster-components/dagster_components/core/schema/metadata.py
index fee035d49e77f..f52ef4bea6052 100644
--- a/python_modules/libraries/dagster-components/dagster_components/core/schema/metadata.py
+++ b/python_modules/libraries/dagster-components/dagster_components/core/schema/metadata.py
@@ -65,7 +65,9 @@ def _subschemas_on_path(
     # List[ComplexType] (e.g.) will contain a reference to the complex type schema in the
     # top-level $defs, so we dereference it here.
     if "$ref" in subschema:
-        subschema = json_schema["$defs"].get(subschema["$ref"][len(REF_BASE) :])
+        # depending on the pydantic version, the extras may be stored with the reference or not
+        extras = {k: v for k, v in subschema.items() if k != "$ref"}
+        subschema = {**json_schema["$defs"].get(subschema["$ref"][len(REF_BASE) :]), **extras}
 
     yield subschema
     if len(valpath) == 0:
diff --git a/python_modules/libraries/dagster-components/dagster_components_tests/rendering_tests/test_schema_resolution.py b/python_modules/libraries/dagster-components/dagster_components_tests/rendering_tests/test_schema_resolution.py
index 4cc192a68d050..aa6a5e40667c8 100644
--- a/python_modules/libraries/dagster-components/dagster_components_tests/rendering_tests/test_schema_resolution.py
+++ b/python_modules/libraries/dagster-components/dagster_components_tests/rendering_tests/test_schema_resolution.py
@@ -70,8 +70,8 @@ def test_allow_render(path, expected: bool) -> None:
         (["inner_seq"], set()),
         (["container_optional_scoped"], {"a", "b"}),
         (["container_optional_scoped", "inner"], {"a", "b"}),
-        # (["container_optional_scoped", "inner_scoped"], {"a", "b", "c", "d"}),
-        # (["container_optional_scoped", "inner_scoped", "a"], {"a", "b", "c", "d"}),
+        (["container_optional_scoped", "inner_scoped"], {"a", "b", "c", "d"}),
+        (["container_optional_scoped", "inner_scoped", "a"], {"a", "b", "c", "d"}),
     ],
 )
 def test_get_available_scope(path, expected: Set[str]) -> None: