fix: rewrite composed refs inside extensions

bundler/bundler_ref_rewrite_test.go

@@ -122,6 +122,108 @@ paths:
 	assertNoFilePathRefs(t, bundled)
 }
 
+func TestBundlerComposed_RewritesRootVendorExtensionRefsToComposedComponents(t *testing.T) {
+	tmpDir := t.TempDir()
+
+	mainSpec := `openapi: 3.1.0
+info:
+  title: Root Vendor Extension Ref Test
+  version: 1.0.0
+tags:
+  - name: things
+    x-related-schemas:
+      - $ref: './models.yaml#/components/schemas/Thing'
+paths:
+  /things:
+    get:
+      tags:
+        - things
+      responses:
+        '200':
+          description: OK
+          content:
+            application/json:
+              schema:
+                $ref: './models.yaml#/components/schemas/Thing'`
+
+	modelsSpec := `components:
+  schemas:
+    Thing:
+      type: object
+      properties:
+        id:
+          type: string`
+
+	require.NoError(t, os.WriteFile(filepath.Join(tmpDir, "openapi.yaml"), []byte(mainSpec), 0644))
+	require.NoError(t, os.WriteFile(filepath.Join(tmpDir, "models.yaml"), []byte(modelsSpec), 0644))
+
+	mainBytes, err := os.ReadFile(filepath.Join(tmpDir, "openapi.yaml"))
+	require.NoError(t, err)
+
+	config := datamodel.NewDocumentConfiguration()
+	config.BasePath = tmpDir
+
+	bundled, err := BundleBytesComposed(mainBytes, config, nil)
+	require.NoError(t, err)
+
+	bundledStr := string(bundled)
+	assert.Contains(t, bundledStr, "x-related-schemas:")
+	assert.Contains(t, bundledStr, `$ref: '#/components/schemas/Thing'`)
+	assert.NotContains(t, bundledStr, "./models.yaml#/components/schemas/Thing")
+	assertNoFilePathRefs(t, bundled)
+}
+
+func TestBundlerComposed_RewritesNestedVendorExtensionRefsToComposedComponents(t *testing.T) {
+	tmpDir := t.TempDir()
+
+	mainSpec := `openapi: 3.1.0
+info:
+  title: Nested Vendor Extension Ref Test
+  version: 1.0.0
+tags:
+  - name: things
+    x-related-field:
+      $ref: './models.yaml#/components/schemas/Thing/properties/id'
+paths:
+  /things:
+    get:
+      tags:
+        - things
+      responses:
+        '200':
+          description: OK
+          content:
+            application/json:
+              schema:
+                $ref: './models.yaml#/components/schemas/Thing'`
+
+	modelsSpec := `components:
+  schemas:
+    Thing:
+      type: object
+      properties:
+        id:
+          type: string`
+
+	require.NoError(t, os.WriteFile(filepath.Join(tmpDir, "openapi.yaml"), []byte(mainSpec), 0644))
+	require.NoError(t, os.WriteFile(filepath.Join(tmpDir, "models.yaml"), []byte(modelsSpec), 0644))
+
+	mainBytes, err := os.ReadFile(filepath.Join(tmpDir, "openapi.yaml"))
+	require.NoError(t, err)
+
+	config := datamodel.NewDocumentConfiguration()
+	config.BasePath = tmpDir
+
+	bundled, err := BundleBytesComposed(mainBytes, config, nil)
+	require.NoError(t, err)
+
+	bundledStr := string(bundled)
+	assert.Contains(t, bundledStr, "x-related-field:")
+	assert.Contains(t, bundledStr, `$ref: '#/components/schemas/Thing/properties/id'`)
+	assert.NotContains(t, bundledStr, "./models.yaml#/components/schemas/Thing/properties/id")
+	assertNoFilePathRefs(t, bundled)
+}
+
 // TestBundlerComposedWithOrigins_TransitiveExternalRefs verifies transitive refs with origin tracking.
 // When schemas use non-standard paths like #/definitions/..., they get inlined rather than composed.
 func TestBundlerComposedWithOrigins_TransitiveExternalRefs(t *testing.T) {

bundler/composer_functions.go

@@ -109,6 +109,41 @@ func processedRefFor(
 	return processedNodes.GetOrZero(fullDefinition)
 }
 
+func composedRefFor(
+	processedNodes *orderedmap.Map[string, *processRef],
+	absoluteKey string,
+) (string, bool) {
+	if processedNodes == nil {
+		return "", false
+	}
+
+	longestKey := ""
+	var longestRef *processRef
+	for key, pr := range processedNodes.FromOldest() {
+		if pr == nil || len(pr.location) == 0 {
+			continue
+		}
+		if key == absoluteKey {
+			continue
+		}
+		if !strings.HasPrefix(absoluteKey, key) {
+			continue
+		}
+		suffix := strings.TrimPrefix(absoluteKey, key)
+		if suffix == "" || !strings.HasPrefix(suffix, "/") {
+			continue
+		}
+		if len(key) > len(longestKey) {
+			longestKey = key
+			longestRef = pr
+		}
+	}
+	if longestRef == nil {
+		return "", false
+	}
+	return "#/" + joinLocationAsJSONPointer(longestRef.location) + strings.TrimPrefix(absoluteKey, longestKey), true
+}
+
 func calculateCollisionName(name, pointer, delimiter string, iteration int) string {
 	jsonPointer := strings.Split(pointer, "#/")
 	if len(jsonPointer) == 2 {
@@ -759,7 +794,7 @@ func walkAndRewriteRefs(
 			// Track extension scope
 			childInExtension := inExtension || strings.HasPrefix(keyNode.Value, "x-")
 
-			if keyNode.Value == "$ref" && valueNode.Kind == yaml.ScalarNode && !inExtension {
+			if keyNode.Value == "$ref" && valueNode.Kind == yaml.ScalarNode {
 				newRef := resolveRefToComposed(valueNode.Value, sourceIdx, processedNodes, rolodex)
 				if newRef != valueNode.Value {
 					valueNode.Value = newRef
@@ -852,6 +887,10 @@ func resolveRefToComposed(
 	// Only rewrite if the target was actually composed into the bundled output.
 	// This prevents dangling refs when SearchIndexForReference resolves something
 	// that never made it into processedNodes.
+	if composedRef, ok := composedRefFor(processedNodes, absoluteKey); ok {
+		return composedRef
+	}
+
 	if processedNodes.GetOrZero(absoluteKey) == nil {
 		return refValue
 	}

bundler/composer_functions_test.go

@@ -274,6 +274,41 @@ func TestProcessedRefFor(t *testing.T) {
 	assert.Nil(t, processedRefFor(processedNodes, "/tmp/missing.yaml#/Thing", nil))
 }
 
+func TestComposedRefFor(t *testing.T) {
+	got, ok := composedRefFor(nil, "/tmp/common.yaml#/components/schemas/Thing/properties/id")
+	assert.False(t, ok)
+	assert.Empty(t, got)
+
+	processedNodes := orderedmap.New[string, *processRef]()
+	processedNodes.Set("/tmp/common.yaml#/components/schemas/Nil", nil)
+	processedNodes.Set("/tmp/common.yaml#/components/schemas/Empty", &processRef{})
+	processedNodes.Set("/tmp/common.yaml#/components/schemas/Thing", &processRef{
+		location: []string{"components", "schemas", "Thing"},
+	})
+	processedNodes.Set("/tmp/common.yaml#/components/schemas/Thing/properties", &processRef{
+		location: []string{"components", "schemas", "ThingProperties"},
+	})
+	processedNodes.Set("/tmp/common.yaml#/components/schemas/Thing/properties/id", &processRef{
+		location: []string{"components", "schemas", "ExactMatch"},
+	})
+
+	got, ok = composedRefFor(processedNodes, "/tmp/common.yaml#/components/schemas/Thing/properties/id/type")
+	require.True(t, ok)
+	assert.Equal(t, "#/components/schemas/ExactMatch/type", got)
+
+	got, ok = composedRefFor(processedNodes, "/tmp/common.yaml#/components/schemas/ThingExtra")
+	assert.False(t, ok)
+	assert.Empty(t, got)
+
+	got, ok = composedRefFor(processedNodes, "/tmp/common.yaml#/components/schemas/Thing/properties/id")
+	assert.True(t, ok)
+	assert.Equal(t, "#/components/schemas/ThingProperties/id", got)
+
+	got, ok = composedRefFor(orderedmap.New[string, *processRef](), "/tmp/common.yaml#/components/schemas/Missing")
+	assert.False(t, ok)
+	assert.Empty(t, got)
+}
+
 func TestInlineProcessRef(t *testing.T) {
 	assert.Nil(t, inlineProcessRef(nil))